bump version

Updated README to include details about v1.5.0 bug fix and new features.

README.md

@@ -7,13 +7,19 @@
 [![Hugging Face](https://img.shields.io/badge/Hugging%20Face-yellow?logo=huggingface&logoColor=white)](https://huggingface.co/charrywhite/LanPaint)
 [![Blog](https://img.shields.io/badge/📝-Blog-9cf)](https://scraed.github.io/scraedBlog/)
 [![GitHub stars](https://img.shields.io/github/stars/scraed/LanPaint)](https://github.com/scraed/LanPaint/stargazers)
-[![Discord](https://img.shields.io/badge/Discord-5865F2?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/aCGZutBV)
+[![Discord](https://img.shields.io/badge/Discord-5865F2?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/yN5wYDE6W4)
 </div>
 
 
 Universally applicable inpainting ability for every model. LanPaint sampler lets the model "think" through multiple iterations before denoising, enabling you to invest more computation time for superior inpainting quality.  
 
-This is the official implementation of ["LanPaint: Training-Free Diffusion Inpainting with Asymptotically Exact and Fast Conditional Sampling"](https://arxiv.org/abs/2502.03491), accepted by TMLR. The repository is for ComfyUI extension. Local Python benchmark code is published here: [LanPaintBench](https://github.com/scraed/LanPaintBench).
+This is the official implementation of ["LanPaint: Training-Free Diffusion Inpainting with Asymptotically Exact and Fast Conditional Sampling"](https://arxiv.org/abs/2502.03491), accepted by TMLR. 
+
+The repository is for ComfyUI extension. 
+
+Diffusers Support: [LanPaint-Diffusers](https://github.com/charrywhite/LanPaint-diffusers) by [@charrywhite](https://github.com/charrywhite/)
+
+Benchmark code for paper reproduce: [LanPaintBench](https://github.com/scraed/LanPaintBench).
 
 ## Citation
 
@@ -31,14 +37,25 @@ note={}
 ```
 **🎉 NEW 2026: Join our discord!**
 
-[Join our Discord](https://discord.gg/aCGZutBV) to share experiences, discuss features, and explore future development.
+[Join our Discord](https://discord.gg/yN5wYDE6W4) to share experiences, discuss features, and explore future development.
+
+`v1.5.0` fixes an important hidden bug that reduced performance and could blur images (especially with `z-image-base`) and also boosts overall LanPaint performance across other models. 
+If your inpainting results have wierd (glowing / broken) mask boundary, check this [issue](https://github.com/scraed/LanPaint/issues/80).
 
 **🎬 NEW: LanPaint now supports inpainting and outpainting based on Z-Image!**
 
+
+
 | Original | Masked | Inpainted |
 |:--------:|:------:|:---------:|
 | ![Original Z-image](https://github.com/scraed/LanPaint/blob/master/examples/Example_21/Original_No_Mask.png) | ![Masked Z-image](https://github.com/scraed/LanPaint/blob/master/examples/Example_21/Masked_Load_Me_in_Loader.png) | ![Inpainted Z-image](https://github.com/scraed/LanPaint/blob/master/examples/Example_21/InPainted_Drag_Me_to_ComfyUI.png) |
 
+**🎬 NEW: LanPaint now supports Z-Image-Base too!**
+
+| Original | Masked | Inpainted |
+|:--------:|:------:|:---------:|
+| ![Original Z-image-base](https://github.com/scraed/LanPaint/blob/master/examples/Example_25/Original_No_Mask.png) | ![Masked Z-image-base](https://github.com/scraed/LanPaint/blob/master/examples/Example_25/Masked_Load_Me_in_Loader.png) | ![Inpainted Z-image-base](https://github.com/scraed/LanPaint/blob/master/examples/Example_25/InPainted_Drag_Me_to_ComfyUI.png) |
+
 
 **🎬 NEW: LanPaint now supports video inpainting and outpainting based on Wan 2.2!**
 
@@ -67,7 +84,9 @@ Check our latest [Wan 2.2 Video Examples](#video-examples-beta), [Wan 2.2 Image
   - [Resource Consumption](#resource-consumption)
 - [Image Examples](#image-examples)
   - [Flux.2.Dev](#example-flux2dev-inpaintlanpaint-k-sampler-5-steps-of-thinking)
+  - [Flux 2 klein](#example-flux-2-klein-inpaintlanpaint-k-sampler-2-steps-of-thinking)
   - [Z-image](#example-z-image-inpaintlanpaint-k-sampler-5-steps-of-thinking)
+  - [Z-image-base](#example-z-image-base-inpaintlanpaint-k-sampler-3-steps-of-thinking)
   - [Hunyuan T2I](#example-hunyuan-t2i-inpaintlanpaint-k-sampler-5-steps-of-thinking)
   - [Wan 2.2 T2I](#example-wan22-inpaintlanpaint-k-sampler-5-steps-of-thinking)
   - [Wan 2.2 T2I with reference](#example-wan22-partial-inpaintlanpaint-k-sampler-5-steps-of-thinking)
@@ -90,7 +109,7 @@ Check our latest [Wan 2.2 Video Examples](#video-examples-beta), [Wan 2.2 Image
 
 ## Features
 
-- **Universal Compatibility** – Works instantly with almost any model (**Z-image, Hunyuan, Wan 2.2, Qwen Image/Edit, HiDream, SD 3.5, Flux-series, SDXL, SD 1.5 or custom LoRAs**) and ControlNet.  
+- **Universal Compatibility** – Works instantly with almost any model (**Z-image, Z-image-base, Hunyuan, Wan 2.2, Qwen Image/Edit, HiDream, SD 3.5, Flux-series, SDXL, SD 1.5 or custom LoRAs**) and ControlNet.  
 ![Inpainting Result 13](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_13.jpg) 
 - **No Training Needed** – Works out of the box with your existing model.  
 - **Easy to Use** – Same workflow as standard ComfyUI KSampler.  
@@ -273,6 +292,24 @@ LanPaint also supports inpainting with the Z-image text-to-image model.
 
 You can download the Z-image model for ComfyUI from [Z-image](https://docs.comfy.org/zh-CN/tutorials/image/z-image/z-image-turbo).
 
+### Example Z-image-base: InPaint(LanPaint K Sampler, 3 steps of thinking)
+LanPaint also supports inpainting with the Z-image-base model.
+
+**Warning (stability)**: Z-image-base can easily diverge with LanPaint. Start with **small `LanPaint_StepSize`** and **fewer thinking iterations** (lower `LanPaint_NumSteps`) and increase gradually only if stable.
+
+<details open>
+<summary>View Original / Masked / Inpainted Comparison</summary>
+
+| Original | Masked | Inpainted |
+|:--------:|:------:|:---------:|
+| ![Original Z-image-base](https://github.com/scraed/LanPaint/blob/master/examples/Example_25/Original_No_Mask.png) | ![Masked Z-image-base](https://github.com/scraed/LanPaint/blob/master/examples/Example_25/Masked_Load_Me_in_Loader.png) | ![Inpainted Z-image-base](https://github.com/scraed/LanPaint/blob/master/examples/Example_25/InPainted_Drag_Me_to_ComfyUI.png) |
+
+</details>
+
+[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_25)
+
+Workflow template (JSON): [Z_image_base_Inpaint.json](https://github.com/scraed/LanPaint/blob/master/example_workflows/Z_image_base_Inpaint.json)
+
 ### Example Wan2.2: Partial InPaint(LanPaint K Sampler, 5 steps of thinking)
 Sometimes we don't want to inpaint completely new content, but rather let the inpainted image reference the original image. One option to achieve this is to inpaint with an edit model like Qwen Image Edit. Another option is to perform a partial inpaint: allowing the diffusion process to start at some middle steps rather than from 0.
 
@@ -342,6 +379,22 @@ You need to follow the ComfyUI version of [SD 3.5 workflow](https://comfyui-wiki
 
 (Note: Prompt First mode is disabled on Flux.2.Dev. As it does not use CFG guidance.)
 
+### Example Flux 2 klein: InPaint(LanPaint K Sampler, 2 steps of thinking)
+
+<details open>
+<summary>View Original / Masked / Inpainted Comparison</summary>
+
+| Original | Masked | Inpainted |
+|:--------:|:------:|:---------:|
+| ![Original Flux 2 klein](https://github.com/scraed/LanPaint/blob/master/examples/Example_24/Original_No_Mask.png) | ![Masked Flux 2 klein](https://github.com/scraed/LanPaint/blob/master/examples/Example_24/Masked_Load_Me_in_Loader.png) | ![Inpainted Flux 2 klein](https://github.com/scraed/LanPaint/blob/master/examples/Example_24/InPainted_Drag_Me_to_ComfyUI.png) |
+
+</details>
+
+[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_24)
+
+[Model Used in This Example](https://docs.comfy.org/zh-CN/tutorials/flux/flux-2-klein). If you have quality problem on Comfy 0.11 and 0.12, check [this issue](https://github.com/scraed/LanPaint/issues/80).
+
+
 ### 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)
@@ -462,6 +515,10 @@ Submit a PR to add your tutorial/video here, or open an [Issue](https://github.c
 [Working togather with crop&stitch](https://github.com/scraed/LanPaint/issues/46)
 
 ## Updates
+- 2026/03/02
+    - `v1.5.0`: Fixed a hidden bug that hurt performance and caused image blur (especially on `z-image-base`), and improved overall LanPaint performance on other models too.
+- 2026/01/30
+    - Add Z-image-base documentation and Example_25 workflow images.
 - 2025/08/08
     - Add Qwen image support
 - 2025/06/21

pyproject.toml

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

src/LanPaint/earlystop.py

@@ -0,0 +1,337 @@
+"""
+Early Stop Logic Contributed by `https://github.com/godnight10061`.
+"""
+
+import inspect
+from typing import Any, Callable, Optional
+
+import torch
+
+from .types import LangevinState
+
+
+def _clamp01(val: float) -> float:
+    if val <= 0.0:
+        return 0.0
+    if val >= 1.0:
+        return 1.0
+    return val
+
+
+def _abt_scale(abt_val: float) -> float:
+    """
+    Smooth, parameter-free scale based on outer-step noise level.
+
+    - 0 at abt=0/1 (disable at extreme noise / extreme tail)
+    - 1 at abt=0.5 (mid-schedule)
+    """
+    abt_val = _clamp01(abt_val)
+    return _clamp01(4.0 * abt_val * (1.0 - abt_val))
+
+
+def _boundary_weight(latent_mask: torch.Tensor, inpaint_weight: torch.Tensor) -> Optional[torch.Tensor]:
+    """
+    Return a 4-neighbor boundary weight: unknown pixels adjacent to known pixels.
+
+    This replaces the previous dilation-based "ring" (kernel/padding) and has no tunable hyperparameters.
+    """
+    if latent_mask.dim() != 4:
+        return None
+
+    known = latent_mask > 0.5
+    neighbor_known = torch.zeros_like(known)
+    neighbor_known[:, :, 1:, :] |= known[:, :, :-1, :]
+    neighbor_known[:, :, :-1, :] |= known[:, :, 1:, :]
+    neighbor_known[:, :, :, 1:] |= known[:, :, :, :-1]
+    neighbor_known[:, :, :, :-1] |= known[:, :, :, 1:]
+
+    boundary = (~known) & neighbor_known
+    return boundary.to(dtype=torch.float32) * inpaint_weight
+
+
+def _weighted_mse(t1: torch.Tensor, t2: torch.Tensor, weight: torch.Tensor) -> float:
+    diff_sq = (t1.to(dtype=torch.float32) - t2.to(dtype=torch.float32)) ** 2
+    denom = torch.sum(weight) + 1e-12
+    return float((torch.sum(diff_sq * weight) / denom).item())
+
+
+class LanPaintEarlyStopper:
+    """
+    Per-step early-stop logic for LanPaint inner (Langevin) iterations.
+    """
+
+    @classmethod
+    def from_options(
+        cls,
+        *,
+        model_options: Optional[dict],
+        latent_mask: torch.Tensor,
+        abt: torch.Tensor,
+        default_threshold: float,
+        default_patience: int,
+        default_distance_fn: Optional[Callable[..., Any]],
+    ) -> Optional["LanPaintEarlyStopper"]:
+        semantic_stop = model_options.get("lanpaint_semantic_stop") if isinstance(model_options, dict) else None
+
+        threshold = float(default_threshold)
+        patience = int(default_patience)
+        distance_fn = default_distance_fn
+        # distance_fn contract: return None (use default metric) or a scalar (Python number / 0-d (1-element) torch.Tensor)
+
+        if isinstance(semantic_stop, dict):
+            threshold = float(semantic_stop.get("threshold", threshold))
+            patience = int(semantic_stop.get("patience", patience))
+            distance_fn = semantic_stop.get("distance_fn", distance_fn)
+
+            # Backward compatibility: map legacy 'min_steps' to a patience floor so it is not an independent knob.
+            if patience > 0:
+                min_steps = semantic_stop.get("min_steps")
+                if min_steps is not None:
+                    try:
+                        min_steps_int = int(min_steps)
+                    except (TypeError, ValueError):
+                        min_steps_int = 0
+                    if min_steps_int > 1:
+                        patience = max(patience, min_steps_int - 1)
+
+        enabled_early_stop = (threshold > 0.0) and (patience > 0)
+        # Require N+1 consecutive stable checks:
+        # - the first stable step sets patience_counter to 1
+        # - `patience=1` therefore stops after 2 stable steps
+        patience_eff = max(1, patience) + 1
+        threshold_eff = threshold
+        inpaint_weight = ring_weight = trace = abt_val = None
+
+        if enabled_early_stop:
+            try:
+                abt_val = float(torch.mean(abt).item())
+            except (TypeError, ValueError):
+                abt_val = 0.0
+
+            threshold_eff = threshold * _abt_scale(abt_val)
+            if threshold_eff <= 0.0:
+                enabled_early_stop = False
+            else:
+                inpaint_weight = (1 - latent_mask).to(dtype=torch.float32)
+                if float(torch.sum(inpaint_weight).item()) < 1e-6:
+                    enabled_early_stop = False
+                else:
+                    ring_weight = _boundary_weight(latent_mask, inpaint_weight)
+                    if isinstance(model_options, dict):
+                        trace = model_options.get("lanpaint_semantic_trace")
+
+        if not enabled_early_stop:
+            return None
+
+        # Pre-fetch trace keys to avoid repeated dict lookups
+        bench_case_id = bench_outer_step = bench_timestep = None
+        if isinstance(trace, list) and isinstance(model_options, dict):
+            bench_case_id = model_options.get("bench_case_id")
+            bench_outer_step = model_options.get("bench_outer_step")
+            bench_timestep = model_options.get("bench_timestep")
+
+        return cls(
+            enabled=enabled_early_stop,
+            threshold=threshold,
+            threshold_eff=threshold_eff,
+            patience_eff=patience_eff,
+            inpaint_weight=inpaint_weight,
+            ring_weight=ring_weight,
+            distance_fn=distance_fn,
+            trace=trace,
+            bench_case_id=bench_case_id,
+            bench_outer_step=bench_outer_step,
+            bench_timestep=bench_timestep,
+            abt_val=abt_val,
+        )
+
+    def __init__(
+        self,
+        *,
+        enabled: bool,
+        threshold: float,
+        threshold_eff: float,
+        patience_eff: int,
+        inpaint_weight: Optional[torch.Tensor],
+        ring_weight: Optional[torch.Tensor],
+        distance_fn: Optional[Callable[..., Any]] = None,
+        trace: Optional[list] = None,
+        bench_case_id: Any = None,
+        bench_outer_step: Any = None,
+        bench_timestep: Any = None,
+        abt_val: Optional[float] = None,
+    ) -> None:
+        self.enabled = bool(enabled)
+        self.threshold = float(threshold)
+        self.threshold_eff = float(threshold_eff)
+        self.patience_eff = int(patience_eff)
+
+        self.inpaint_weight = inpaint_weight
+        self.ring_weight = ring_weight
+
+        self.trace = trace
+        self.bench_case_id = bench_case_id
+        self.bench_outer_step = bench_outer_step
+        self.bench_timestep = bench_timestep
+        self.abt_val = abt_val
+
+        self.patience_counter = 0
+        self.x0_anchor = None
+
+        self._dist_wrapper = self._wrap_distance_fn(distance_fn) if self.enabled else None
+
+    @property
+    def has_custom_distance_fn(self) -> bool:
+        return self._dist_wrapper is not None
+
+    @staticmethod
+    def _wrap_distance_fn(distance_fn: Optional[Callable[..., Any]]):
+        """
+        Wrap a user-provided `distance_fn` into a normalized callable: fn(prev, cur, ctx) -> dist|None.
+
+        Supported signatures:
+        - 3+ positional (or *args): `distance_fn(prev, cur, ctx)`
+        - explicit / **kwargs ctx: `distance_fn(prev, cur, ctx=ctx)`
+        - default 2-arg: `distance_fn(cur, prev)`
+
+        Return contract: None (use default metric) or a scalar (Python number / 0-d (1-element) torch.Tensor).
+        """
+        if not callable(distance_fn):
+            return None
+
+        try:
+            sig = inspect.signature(distance_fn)
+            params = list(sig.parameters.values())
+
+            has_ctx_param = "ctx" in sig.parameters
+            has_var_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params)
+            has_var_pos = any(p.kind == inspect.Parameter.VAR_POSITIONAL for p in params)
+
+            pos_params = [
+                p
+                for p in params
+                if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+            ]
+
+            if len(pos_params) >= 3 or has_var_pos:
+                # 3-arg positional: fn(prev, cur, ctx)
+                return lambda p, c, ctx: distance_fn(p, c, ctx)
+            if has_ctx_param or has_var_kw:
+                # keyword ctx: fn(prev, cur, ctx=ctx)
+                return lambda p, c, ctx: distance_fn(p, c, ctx=ctx)
+
+            # Default 2-arg: fn(cur, prev)
+            return lambda p, c, ctx: distance_fn(c, p)
+        except (ValueError, TypeError):
+            # Fallback for built-ins or complex callables.
+            def fallback_wrapper(p, c, ctx):
+                try:
+                    return distance_fn(p, c, ctx)
+                except TypeError as e:
+                    tb = e.__traceback__
+                    if tb is not None and tb.tb_frame.f_code is not fallback_wrapper.__code__:
+                        raise
+                    return distance_fn(c, p)
+
+            return fallback_wrapper
+
+    def step(
+        self,
+        *,
+        i: int,
+        n_steps: int,
+        x_t_before: torch.Tensor,
+        x_t_after: torch.Tensor,
+        x_t_prev_for_custom: Optional[torch.Tensor],
+        prev_args: Any,
+        args: Any,
+        ctx: dict,
+    ) -> bool:
+        if not self.enabled:
+            return False
+
+        # 'inpaint_weight' is guaranteed to be set when enabled is True in the caller.
+        inpaint = self.inpaint_weight
+        if inpaint is None:
+            return False
+
+        dist = None
+        custom_dist = False
+        dist_inpaint = dist_ring = dist_drift = x0_prev = x0_cur = None
+
+        if self._dist_wrapper is not None:
+            dist = self._dist_wrapper(x_t_prev_for_custom, x_t_after, ctx)
+            if dist is not None:
+                if isinstance(dist, torch.Tensor):
+                    if dist.numel() != 1:
+                        raise TypeError("distance_fn must return None or a scalar / 0-d (1-element) tensor")
+                    dist = float(dist.item())
+                else:
+                    dist = float(dist)
+        custom_dist = dist is not None
+
+        if dist is None:
+            def _get_x0(arg: Any) -> Optional[torch.Tensor]:
+                if isinstance(arg, LangevinState):
+                    return arg.x0
+                if isinstance(arg, tuple) and len(arg) >= 3:
+                    return arg[2]
+                return None
+
+            x0_prev = _get_x0(prev_args)
+            x0_cur = _get_x0(args)
+
+            if x0_prev is not None and x0_cur is not None:
+                dist_inpaint = _weighted_mse(x0_cur, x0_prev, inpaint)
+                dist_ring = _weighted_mse(x0_cur, x0_prev, self.ring_weight) if self.ring_weight is not None else None
+                dist = dist_inpaint if dist_ring is None else max(dist_inpaint, dist_ring)
+            else:
+                dist_inpaint = _weighted_mse(x_t_after, x_t_before, inpaint)
+                dist = dist_inpaint
+
+        threshold_used = self.threshold if custom_dist else self.threshold_eff
+
+        # Drift guard (only for default metric with x0_cur).
+        if x0_cur is not None and not custom_dist:
+            if dist <= threshold_used:
+                if self.x0_anchor is None:
+                    self.x0_anchor = x0_cur.detach()
+                else:
+                    drift_inpaint = _weighted_mse(x0_cur, self.x0_anchor, inpaint)
+                    drift_ring = _weighted_mse(x0_cur, self.x0_anchor, self.ring_weight) if self.ring_weight is not None else None
+                    dist_drift = drift_inpaint if drift_ring is None else max(drift_inpaint, drift_ring)
+                    dist = max(dist, dist_drift)
+            else:
+                self.x0_anchor = None
+
+        if dist <= threshold_used:
+            self.patience_counter += 1
+        else:
+            self.patience_counter = 0
+            self.x0_anchor = None
+
+        should_stop = self.patience_counter >= self.patience_eff
+
+        if isinstance(self.trace, list):
+            self.trace.append(
+                {
+                    "case_id": self.bench_case_id,
+                    "outer_step": self.bench_outer_step,
+                    "bench_timestep": self.bench_timestep,
+                    "inner_step": i + 1,
+                    "dist": dist,
+                    "dist_inpaint": None if dist_inpaint is None else float(dist_inpaint),
+                    "dist_ring": None if dist_ring is None else float(dist_ring),
+                    "dist_drift": None if dist_drift is None else float(dist_drift),
+                    "threshold": float(threshold_used),
+                    "threshold_eff": float(self.threshold_eff),
+                    "patience_counter": int(self.patience_counter),
+                    "patience_eff": int(self.patience_eff),
+                    "abt": None if self.abt_val is None else float(self.abt_val),
+                    "custom_dist": bool(custom_dist),
+                    "stopped": bool(should_stop),
+                }
+            )
+
+        return bool(should_stop)
+

src/LanPaint/lanpaint.py

@@ -1,9 +1,11 @@
 import torch
 from .utils import StochasticHarmonicOscillator
 from functools import partial
+from .earlystop import LanPaintEarlyStopper
+from .types import LangevinState
 
 class LanPaint():
-    def __init__(self, Model, NSteps, Friction, Lambda, Beta, StepSize, IS_FLUX = False, IS_FLOW = False):
+    def __init__(self, Model, NSteps, Friction, Lambda, Beta, StepSize, IS_FLUX = False, IS_FLOW = False, EarlyStopThreshold = 0.0, EarlyStopPatience = 1, EarlyStopHook = None):
         self.n_steps = NSteps
         self.chara_lamb = Lambda
         self.IS_FLUX = IS_FLUX
@@ -13,6 +15,9 @@ class LanPaint():
         self.friction = Friction
         self.chara_beta = Beta
         self.img_dim_size = None
+        self.early_stop_threshold = EarlyStopThreshold
+        self.early_stop_patience = EarlyStopPatience
+        self.early_stop_hook = EarlyStopHook
 
     def add_none_dims(self, array):
         # Create a tuple with ':' for the first dimension and 'None' repeated num_nones times
@@ -22,6 +27,16 @@ class LanPaint():
         # Create a tuple with ':' for the first dimension and 'None' repeated num_nones times
         index = (slice(None),) + (0,) * (self.img_dim_size-1)
         return array[index]
+    def unpack_model_output(self, output):
+        # Some guider/model wrappers return one denoised latent, others return
+        # both the normal and BIG-guidance denoised latents.
+        if isinstance(output, (tuple, list)):
+            if len(output) >= 2:
+                return output[0], output[1]
+            if len(output) == 1:
+                return output[0], output[0]
+            raise ValueError("Model output is empty")
+        return output, output
     def __call__(self, x, latent_image, noise, sigma, latent_mask, current_times, model_options, seed, n_steps=None):
         self.img_dim_size = len(x.shape)
         self.latent_image = latent_image
@@ -32,6 +47,7 @@ class LanPaint():
             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):
+        input_x = x
         VE_Sigma, abt, Flow_t = current_times
 
         step_size = self.step_size * (1 - abt)
@@ -51,17 +67,59 @@ class LanPaint():
         ############ 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
+        stopper = LanPaintEarlyStopper.from_options(
+            model_options=model_options if isinstance(model_options, dict) else None,
+            latent_mask=latent_mask,
+            abt=abt,
+            default_threshold=self.early_stop_threshold,
+            default_patience=self.early_stop_patience,
+            default_distance_fn=self.early_stop_hook,
+        )
+
         for i in range(n_steps):
             score_func = partial( self.score_model, y = self.latent_image, mask = latent_mask, abt = self.add_none_dims(abt), sigma = self.add_none_dims(VE_Sigma), tflow = self.add_none_dims(Flow_t), model_options = model_options, seed = seed )
+
+            prev_args = args
+            x_t_prev = x_t.detach() if (stopper is not None and stopper.has_custom_distance_fn) else None
+            x_t_before = x_t if (stopper is not None and stopper.enabled) else None
+
             x_t, args = self.langevin_dynamics(x_t, score_func , latent_mask, step_size , current_times, sigma_x = self.add_none_dims(self.sigma_x(abt)), sigma_y = self.add_none_dims(self.sigma_y(abt)), args = args)  
+
+            if stopper is not None:
+                ctx = {
+                    "step": i,
+                    "steps_done": i + 1,
+                    "n_steps": n_steps,
+                    "mask": latent_mask,
+                    "latent_image": self.latent_image,
+                    "current_times": current_times,
+                    "seed": seed,
+                }
+                if stopper.step(
+                    i=i,
+                    n_steps=n_steps,
+                    x_t_before=x_t_before,
+                    x_t_after=x_t,
+                    x_t_prev_for_custom=x_t_prev,
+                    prev_args=prev_args,
+                    args=args,
+                    ctx=ctx,
+                ):
+                    break
+
         if IS_FLUX or IS_FLOW:
             x = x_t / ( self.add_none_dims(abt)**0.5 + (1-self.add_none_dims(abt))**0.5 )
         else:
             x = x_t * ( 1+self.add_none_dims(VE_Sigma)**2 )**0.5 # switch to variance perserving x_t values
         ############ LanPaint Iterations End ###############
         # out is x_0
-        out, _ = self.inner_model(x, sigma, model_options=model_options, seed=seed)
+
+        out, _ = self.unpack_model_output(
+            self.inner_model(x, sigma, model_options=model_options, seed=seed)
+        )
         out = out * (1-latent_mask) + self.latent_image * latent_mask
+
+        input_x.copy_(x)
         return out
 
     def score_model(self, x_t, y, mask, abt, sigma, tflow, model_options, seed):
@@ -69,10 +127,14 @@ class LanPaint():
         if self.IS_FLUX or self.IS_FLOW:
             # compute t for flow model, with a small epsilon compensating for numerical error.
             x = x_t / ( abt**0.5 + (1-abt)**0.5 ) # switch to Gaussian flow matching
-            x_0, x_0_BIG = self.inner_model(x, self.remove_none_dims(tflow), model_options=model_options, seed=seed)
+            x_0, x_0_BIG = self.unpack_model_output(
+                self.inner_model(x, self.remove_none_dims(tflow), model_options=model_options, seed=seed)
+            )
         else:
             x = x_t * ( 1+sigma**2 )**0.5 # switch to variance exploding
-            x_0, x_0_BIG = self.inner_model(x, self.remove_none_dims(sigma), model_options=model_options, seed=seed)
+            x_0, x_0_BIG = self.unpack_model_output(
+                self.inner_model(x, self.remove_none_dims(sigma), model_options=model_options, seed=seed)
+            )
 
         score_x = -(x_t - x_0)
         score_y =  - (1 + lamb) * ( x_t - y )  + lamb * (x_t - x_0_BIG)  
@@ -85,6 +147,13 @@ class LanPaint():
         return beta
 
     def langevin_dynamics(self, x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None):
+        if args is not None and not isinstance(args, LangevinState):
+            if isinstance(args, tuple):
+                if len(args) == 2:
+                    # Backwards compat: older state was (v, C) without x0.
+                    args = LangevinState(args[0], args[1], None)
+                elif len(args) >= 3:
+                    args = LangevinState(args[0], args[1], args[2])
         # 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)
@@ -102,11 +171,10 @@ class LanPaint():
         dt = dtx * (1-mask) + dty * mask
         Gamma = Gamma_x * (1-mask) + Gamma_y * mask
 
-
         def Coef_C(x_t):
-            x0 = self.x0_evalutation(x_t, score, sigma, args)
+            x0 = x_t + score(x_t)
             C = (abt**0.5 * x0  - x_t )/ (1-abt) + A * x_t
-            return C
+            return C, x0
         def advance_time(x_t, v, dt, Gamma, A, C, D):
             dtype = x_t.dtype
             with torch.autocast(device_type=x_t.device.type, dtype=torch.float32):
@@ -143,44 +211,46 @@ class LanPaint():
         def run_damped(x_t, args):
             if args is None:
                 v = None
-                C = Coef_C(x_t)
+                C, x0 = Coef_C(x_t)
                 x_t, v = advance_time(x_t, v, dt, Gamma, A, C, D)
             else:
-                v, C = args
+                v = args.v
+                C = args.C
                 x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D)
-                C_new = Coef_C(x_t)
+                C_new, x0 = Coef_C(x_t)
                 v = v + Gamma**0.5 * ( C_new - C) *dt
                 x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D)
                 C = C_new
-            return x_t, (v, C)
+            # args is (v, C, x0) for the next inner step.
+            return x_t, LangevinState(v, C, x0)
 
         def run_overdamped(x_t, args):
             if args is None:
-                C = Coef_C(x_t)
+                C, x0 = Coef_C(x_t)
                 x_t = advance_time_overdamped(x_t, dt, A, C, D)
             else:
-                _, C = args
+                C = args.C
                 x_t = advance_time_overdamped(x_t, dt / 2, A, C, D)
-                C_new = Coef_C(x_t)
+                C_new, x0 = Coef_C(x_t)
                 x_t = x_t + (C_new - C) * dt
                 x_t = advance_time_overdamped(x_t, dt / 2, A, C, D)
                 C = C_new
-            return x_t, (None, C)
+            # args is (v, C, x0); v is None in the overdamped fallback.
+            return x_t, LangevinState(None, C, x0)
 
         try:
-            x_t_next, (v_next, C_next) = run_damped(x_t, args)
+            x_t_next, state = run_damped(x_t, args)
 
-            if torch.isnan(x_t_next).any() or torch.isnan(v_next).any():
+            v_next = state.v
+            if torch.isnan(x_t_next).any() or (v_next is not None and torch.isnan(v_next).any()):
                 raise ValueError("NaN detected")
 
             x_t = x_t_next
-            v = v_next
-            C = C_next
-
         except Exception:
-            x_t, (v, C) = run_overdamped(x_t, args)
+            x_t, state = run_overdamped(x_t, args)
 
-        return x_t, (v, C)
+        # args is (v, C, x0); v can be None if we fell back to the overdamped update.
+        return x_t, state
 
     def prepare_step_size(self, current_times, step_size, sigma_x, sigma_y):
         # -------------------------------------------------------------------------
@@ -216,9 +286,3 @@ class LanPaint():
         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

@@ -12,6 +12,11 @@ from .lanpaint import LanPaint
 from comfy.model_base import WAN22
 import comfyui_version 
 
+def _version_tuple(value):
+    return tuple(int(part) if part.isdigit() else 0 for part in value.split("."))
+
+COMFYUI_VERSION_060_OR_NEWER = _version_tuple(comfyui_version.__version__) >= (0, 6, 0)
+
 def reshape_mask(input_mask, output_shape,video_inpainting=False):
     dims = len(output_shape) - 2
     print('output shape',output_shape)
@@ -28,7 +33,7 @@ def reshape_mask(input_mask, output_shape,video_inpainting=False):
 
     # Handle 5D output shape (B, C, F, H, W) by ensuring input is 5D
     if len(output_shape) == 5 and input_mask.ndim == 4:
-        if comfyui_version.__version__ >= "0.6.0":
+        if COMFYUI_VERSION_060_OR_NEWER:
             input_mask = input_mask.unsqueeze(2)  # (B, C, 1, H, W)
 
     # Handle video case with temporal dimension
@@ -41,7 +46,7 @@ def reshape_mask(input_mask, output_shape,video_inpainting=False):
         # First reshape input_mask to have proper dimensions for video processing
         # Assume input is (frames, channels, height, width) -> (1, channels, frames, height, width)
         ## if comfy version < 0.6.0
-        if comfyui_version.__version__ < "0.6.0":
+        if not COMFYUI_VERSION_060_OR_NEWER:
             input_mask = input_mask.permute(1, 0, 2, 3).unsqueeze(0)
         print('Video case - input_mask after reshaping:', input_mask.shape)
         # Ensure we have the correct 5D shape: (batch, channels, frames, height, width)
@@ -65,7 +70,7 @@ def reshape_mask(input_mask, output_shape,video_inpainting=False):
         # Handle batch dimension
         mask = repeat_to_batch_size(mask, output_shape[0])
     else:  # Original 2D image case
-        if comfyui_version.__version__ < "0.6.0":
+        if not COMFYUI_VERSION_060_OR_NEWER:
             mask = torch.nn.functional.interpolate(input_mask, size=output_shape[-2:], mode=scale_mode)
         else:
             mask = torch.nn.functional.interpolate(input_mask, size=output_shape[2:], mode=scale_mode)
@@ -222,7 +227,10 @@ class KSAMPLER(comfy.samplers.KSAMPLER):
                                        model_wrap.model_patcher.LanPaint_Beta,
                                        model_wrap.model_patcher.LanPaint_StepSize, 
                                        IS_FLUX = IS_FLUX, 
-                                       IS_FLOW = IS_FLOW)
+                                       IS_FLOW = IS_FLOW,
+                                       EarlyStopThreshold = getattr(model_wrap.model_patcher, "LanPaint_InnerThreshold", 0.0),
+                                       EarlyStopPatience = getattr(model_wrap.model_patcher, "LanPaint_InnerPatience", 1),
+                                       EarlyStopHook = extra_args.get("model_options", {}).get("lanpaint_semantic_hook", None))
         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.
@@ -324,6 +332,8 @@ class LanPaint_KSampler():
         model.LanPaint_NumSteps = LanPaint_NumSteps
         model.LanPaint_Friction = 15.
         model.LanPaint_EarlyStop = 1
+        model.LanPaint_InnerThreshold = 0.0
+        model.LanPaint_InnerPatience = 1
         if LanPaint_PromptMode == "Image First":
             model.LanPaint_cfg_BIG = cfg
         else:
@@ -363,6 +373,8 @@ class LanPaint_KSamplerAdvanced:
                 "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}),
                 "Inpainting_mode": (["🖼️ Image Inpainting", "🎬 Video Inpainting"], {"default": "🖼️ Image Inpainting", "tooltip": "Choose Image mode for photos or Video mode for video frames with temporal consistency"}),
+                "LanPaint_InnerThreshold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.0001, "round": 0.0001, "tooltip": "Early stop threshold for Langevin iterations based on semantic distance. 0.0 to disable. (Contributed by godnight10061)"}),
+                "LanPaint_InnerPatience": ("INT", {"default": 1, "min": 1, "max": 100, "tooltip": "Number of consecutive steps below threshold required to stop. (Contributed by godnight10061)"}),
                      },
                 }
 
@@ -371,7 +383,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, LanPaint_NumSteps=5, LanPaint_Lambda=16.0, LanPaint_StepSize=0.2, LanPaint_Beta=1.0, LanPaint_Friction=15.0, LanPaint_PromptMode="Image First", LanPaint_EarlyStop=1, LanPaint_Info="", Inpainting_mode="🖼️ Image Inpainting"):
+    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, LanPaint_NumSteps=5, LanPaint_Lambda=16.0, LanPaint_StepSize=0.2, LanPaint_Beta=1.0, LanPaint_Friction=15.0, LanPaint_PromptMode="Image First", LanPaint_EarlyStop=1, LanPaint_Info="", Inpainting_mode="🖼️ Image Inpainting", LanPaint_InnerThreshold=0.0, LanPaint_InnerPatience=1):
         force_full_denoise = True
         if return_with_leftover_noise == "enable":
             force_full_denoise = False
@@ -384,6 +396,8 @@ class LanPaint_KSamplerAdvanced:
         model.LanPaint_NumSteps = LanPaint_NumSteps
         model.LanPaint_Friction = LanPaint_Friction
         model.LanPaint_EarlyStop = LanPaint_EarlyStop
+        model.LanPaint_InnerThreshold = LanPaint_InnerThreshold
+        model.LanPaint_InnerPatience = LanPaint_InnerPatience
         if LanPaint_PromptMode == "Image First":
             model.LanPaint_cfg_BIG = cfg
         else:
@@ -505,6 +519,8 @@ class LanPaint_SamplerCustom:
         model.LanPaint_NumSteps = LanPaint_NumSteps
         model.LanPaint_Friction = 15.
         model.LanPaint_EarlyStop = 1
+        model.LanPaint_InnerThreshold = 0.0
+        model.LanPaint_InnerPatience = 1
         if LanPaint_PromptMode == "Image First":
             model.LanPaint_cfg_BIG = cfg
         else:
@@ -556,6 +572,8 @@ class LanPaint_SamplerCustomAdvanced:
                      "LanPaint_PromptMode": (["Image First", "Prompt First"], {"tooltip": "Image First: prioritizes image quality; Prompt First: prioritizes prompt adherence."}),
                      "LanPaint_EarlyStop": ("INT", {"default": 1, "min": 0, "max": 10000, "tooltip": "Steps to stop LanPaint early, preventing irregular patterns."}),
                      "LanPaint_Info": ("STRING", {"default": "LanPaint Custom Sampler Adv. For more info, visit https://github.com/scraed/LanPaint. If you find it useful, please give a star ⭐️!", "multiline": True}),
+                     "LanPaint_InnerThreshold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.0001, "round": 0.0001, "tooltip": "Early stop threshold for Langevin iterations based on semantic distance. 0.0 to disable. (Contributed by godnight10061)"}),
+                     "LanPaint_InnerPatience": ("INT", {"default": 1, "min": 1, "max": 100, "tooltip": "Number of consecutive steps below threshold required to stop. (Contributed by godnight10061)"}),
                     }
                }
 
@@ -566,7 +584,7 @@ class LanPaint_SamplerCustomAdvanced:
 
     CATEGORY = "sampling/custom_sampling"
 
-    def sample(self, noise, guider, sampler, sigmas, latent_image, LanPaint_NumSteps, LanPaint_Lambda, LanPaint_StepSize, LanPaint_Beta, LanPaint_Friction, LanPaint_PromptMode, LanPaint_EarlyStop, LanPaint_Info=""):
+    def sample(self, noise, guider, sampler, sigmas, latent_image, LanPaint_NumSteps, LanPaint_Lambda, LanPaint_StepSize, LanPaint_Beta, LanPaint_Friction, LanPaint_PromptMode, LanPaint_EarlyStop, LanPaint_Info="", LanPaint_InnerThreshold=0.0, LanPaint_InnerPatience=1):
         model = guider.model_patcher
         model.LanPaint_StepSize = LanPaint_StepSize
         model.LanPaint_Lambda = LanPaint_Lambda
@@ -574,6 +592,8 @@ class LanPaint_SamplerCustomAdvanced:
         model.LanPaint_NumSteps = LanPaint_NumSteps
         model.LanPaint_Friction = LanPaint_Friction
         model.LanPaint_EarlyStop = LanPaint_EarlyStop
+        model.LanPaint_InnerThreshold = LanPaint_InnerThreshold
+        model.LanPaint_InnerPatience = LanPaint_InnerPatience
         if LanPaint_PromptMode == "Image First":
             model.LanPaint_cfg_BIG = guider.cfg
         else:

src/LanPaint/types.py

@@ -0,0 +1,10 @@
+from typing import NamedTuple, Optional
+
+import torch
+
+
+class LangevinState(NamedTuple):
+    v: Optional[torch.Tensor]
+    C: Optional[torch.Tensor]
+    x0: Optional[torch.Tensor]
+

tests/test_lanpaint_semantic_stop.py

@@ -0,0 +1,104 @@
+import torch
+
+from src.LanPaint.lanpaint import LanPaint as LanPaintEngine
+
+
+class _DummySampling:
+    def noise_scaling(self, sigma, noise, latent_image):  # type: ignore[no-untyped-def]
+        return latent_image + noise * sigma
+
+
+class _DummyModel:
+    def __init__(self) -> None:
+        self.inner_model = self
+        self.model_sampling = _DummySampling()
+
+    def __call__(self, x, sigma, model_options=None, seed=None):  # type: ignore[no-untyped-def]
+        return x, x
+
+
+def _inputs():  # type: ignore[no-untyped-def]
+    x = torch.zeros((1, 4, 8, 8))
+    latent_image = torch.zeros_like(x)
+    noise = torch.ones_like(x)
+    sigma = torch.tensor([1.0])
+
+    latent_mask = torch.zeros_like(x)
+    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
+    return x, latent_image, noise, sigma, latent_mask, current_times
+
+
+def test_default_semantic_stop_triggers_at_patience_without_custom_distance_fn() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    calls = {"langevin": 0, "with_score": 0, "without_score": 0}
+
+    def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None):  # type: ignore[no-untyped-def]
+        calls["langevin"] += 1
+        if score is None:
+            calls["without_score"] += 1
+        else:
+            calls["with_score"] += 1
+        return x_t, args
+
+    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "patience": 2,
+        }
+    }
+
+    x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
+    engine(x, latent_image, noise, sigma, latent_mask, current_times, model_options=model_options, seed=0, n_steps=10)
+
+    assert calls["langevin"] == 3
+    assert calls["with_score"] == 3
+    assert calls["without_score"] == 0
+
+
+def test_semantic_stop_is_disabled_when_no_inpaint_region() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    calls = {"langevin": 0, "with_score": 0, "without_score": 0}
+
+    def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None):  # type: ignore[no-untyped-def]
+        calls["langevin"] += 1
+        if score is None:
+            calls["without_score"] += 1
+        else:
+            calls["with_score"] += 1
+        return x_t, args
+
+    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "patience": 1,
+        }
+    }
+
+    x, latent_image, noise, sigma, latent_mask, _ = _inputs()
+    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
+    no_inpaint_mask = torch.ones_like(latent_mask)
+    engine(x, latent_image, noise, sigma, no_inpaint_mask, current_times, model_options=model_options, seed=0, n_steps=10)
+
+    assert calls["langevin"] == 10
+    assert calls["with_score"] == 10
+    assert calls["without_score"] == 0

tests/test_sho_regression.py

@@ -1,6 +1,8 @@
 import torch
 from unittest.mock import MagicMock, patch
 from src.LanPaint.lanpaint import LanPaint
+
+
 def test_langevin_dynamics_fallback_on_nan() -> None:
     """Test that langevin_dynamics falls back to overdamped dynamics if damped dynamics produces NaNs."""
     torch.manual_seed(0)
@@ -26,26 +28,18 @@ def test_langevin_dynamics_fallback_on_nan() -> None:
         mock_instance.dynamics.return_value = (nan_tensor, nan_tensor)
         # Execute langevin_dynamics
         # This should try run_damped -> get NaNs -> raise ValueError -> catch -> run_overdamped
-        x_out, (v_out, C_out) = lp.langevin_dynamics(x_t, score, mask, step_size, current_times, sigma_y=1.0)
+        x_out, args_out = lp.langevin_dynamics(x_t, score, mask, step_size, current_times, sigma_y=1.0)
+        assert hasattr(args_out, "v")
+        assert hasattr(args_out, "C")
+        assert hasattr(args_out, "x0")
+        assert args_out[0] is args_out.v
+        assert args_out[1] is args_out.C
+        assert args_out[2] is args_out.x0
+        v_out = args_out[0]
         # Verify that SHO was initialized and dynamics called
         MockSHO.assert_called()
         mock_instance.dynamics.assert_called()
         # Verify result is finite (indicating fallback to overdamped logic was successful)
         assert torch.isfinite(x_out).all(), "Output contains NaNs, fallback failed"
         assert v_out is None or torch.isfinite(v_out).all()
-def test_langevin_dynamics_fallback_on_exception() -> None:
-    """Test that langevin_dynamics falls back to overdamped dynamics if damped dynamics raises Exception."""
-    torch.manual_seed(0)
-    lp = LanPaint(Model=MagicMock(), NSteps=10, Friction=1.0, Lambda=1.0, Beta=1.0, StepSize=0.1)
-    x_t = torch.randn(1, 4, 8, 8)
-    lp.img_dim_size = 4
-    mask = torch.zeros_like(x_t)
-    score = lambda x: torch.zeros_like(x)
-    step_size = torch.tensor([0.1])
-    current_times = (torch.tensor([0.5]), torch.tensor([0.5]), torch.tensor([0.5]))
-    with patch("src.LanPaint.lanpaint.StochasticHarmonicOscillator") as MockSHO:
-        mock_instance = MockSHO.return_value
-        # Configure dynamics to raise Exception
-        mock_instance.dynamics.side_effect = RuntimeError("Simulation exploded")
-        x_out, (v_out, C_out) = lp.langevin_dynamics(x_t, score, mask, step_size, current_times, sigma_y=1.0)
-        assert torch.isfinite(x_out).all()
+