From 060fcc44750dec96c0a8a89e539579d69375a8c2 Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Sun, 11 Jan 2026 23:11:39 +0800
Subject: [PATCH 01/33] Add semantic early-stop for Langevin iterations

---
 src/LanPaint/lanpaint.py             | 59 ++++++++++++++++++++-
 src/LanPaint/nodes.py                | 21 ++++++--
 tests/test_lanpaint_semantic_stop.py | 77 ++++++++++++++++++++++++++++
 3 files changed, 153 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_lanpaint_semantic_stop.py

diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index 10a6344..7fe94ed 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -3,7 +3,7 @@
 from functools import partial
 
 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 +13,9 @@ def __init__(self, Model, NSteps, Friction, Lambda, Beta, StepSize, IS_FLUX = Fa
         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
@@ -51,9 +54,63 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
         ############ 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
+        semantic_stop = None
+        if isinstance(model_options, dict):
+            semantic_stop = model_options.get("lanpaint_semantic_stop")
+
+        threshold = float(self.early_stop_threshold)
+        patience = int(self.early_stop_patience)
+        min_steps = 1
+        distance_fn = self.early_stop_hook
+
+        if isinstance(semantic_stop, dict):
+            threshold = float(semantic_stop.get("threshold", threshold))
+            patience = int(semantic_stop.get("patience", patience))
+            min_steps = int(semantic_stop.get("min_steps", min_steps))
+            distance_fn = semantic_stop.get("distance_fn", distance_fn)
+
+        enabled_early_stop = (threshold > 0.0) and (patience > 0)
+        min_steps = max(1, min_steps)
+        patience = max(1, patience)
+        patience_counter = 0
+
         for i in range(n_steps):
             score_func = partial( self.score_model, y = self.latent_image, mask = latent_mask, abt = self.add_none_dims(abt), sigma = self.add_none_dims(VE_Sigma), tflow = self.add_none_dims(Flow_t), model_options = model_options, seed = seed )
+
+            x_t_prev = x_t.detach().clone() if enabled_early_stop 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 enabled_early_stop and x_t_prev 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,
+                }
+
+                dist = None
+                if callable(distance_fn):
+                    try:
+                        dist = distance_fn(x_t_prev, x_t, ctx)
+                    except TypeError:
+                        dist = distance_fn(x_t, x_t_prev)
+
+                if dist is None:
+                    diff = (x_t - x_t_prev) * (1 - latent_mask)
+                    dist = torch.mean(diff**2).item()
+
+                if float(dist) <= threshold:
+                    patience_counter += 1
+                else:
+                    patience_counter = 0
+
+                if (i + 1) >= min_steps and patience_counter >= patience:
+                    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:
diff --git a/src/LanPaint/nodes.py b/src/LanPaint/nodes.py
index 0cff2af..8efdc19 100644
--- a/src/LanPaint/nodes.py
+++ b/src/LanPaint/nodes.py
@@ -227,7 +227,10 @@ def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=N
                                        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.
@@ -329,6 +332,8 @@ def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, neg
         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:
@@ -366,6 +371,8 @@ def INPUT_TYPES(s):
                 "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_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."}),
+                "LanPaint_InnerPatience": ("INT", {"default": 1, "min": 1, "max": 100, "tooltip": "Number of consecutive steps below threshold required to stop."}),
                 "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"}),
                      },
@@ -376,7 +383,7 @@ def INPUT_TYPES(s):
 
     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_InnerThreshold=0.0, LanPaint_InnerPatience=1, LanPaint_Info="", Inpainting_mode="🖼️ Image Inpainting"):
         force_full_denoise = True
         if return_with_leftover_noise == "enable":
             force_full_denoise = False
@@ -389,6 +396,8 @@ def sample(self, model, add_noise, noise_seed, steps, cfg, sampler_name, schedul
         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:
@@ -510,6 +519,8 @@ def sample(self, model, sampler, sigmas, add_noise, noise_seed, cfg, positive, n
         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:
@@ -560,6 +571,8 @@ def INPUT_TYPES(s):
                      "LanPaint_Friction": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 50.0, "step": 0.1, "tooltip": "Friction parameter for fast Langevin. Lower values speed convergence but may be unstable."}),
                      "LanPaint_PromptMode": (["Image First", "Prompt First"], {"tooltip": "Image First: prioritizes image quality; Prompt First: prioritizes prompt adherence."}),
                      "LanPaint_EarlyStop": ("INT", {"default": 1, "min": 0, "max": 10000, "tooltip": "Steps to stop LanPaint early, preventing irregular patterns."}),
+                     "LanPaint_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."}),
+                     "LanPaint_InnerPatience": ("INT", {"default": 1, "min": 1, "max": 100, "tooltip": "Number of consecutive steps below threshold required to stop."}),
                      "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}),
                     }
                }
@@ -571,7 +584,7 @@ def INPUT_TYPES(s):
 
     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_InnerThreshold, LanPaint_InnerPatience, LanPaint_Info=""):
         model = guider.model_patcher
         model.LanPaint_StepSize = LanPaint_StepSize
         model.LanPaint_Lambda = LanPaint_Lambda
@@ -579,6 +592,8 @@ def sample(self, noise, guider, sampler, sigmas, latent_image, LanPaint_NumSteps
         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:
diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
new file mode 100644
index 0000000..5ed90d2
--- /dev/null
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -0,0 +1,77 @@
+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_semantic_stop_triggers_deterministically_at_min_steps() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    calls = {"langevin": 0, "distance": 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
+        return x_t, args
+
+    def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
+        assert ctx["step"] == calls["distance"]
+        calls["distance"] += 1
+        return 0.0
+
+    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "min_steps": 3,
+            "patience": 1,
+            "distance_fn": distance_fn,
+        }
+    }
+
+    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["distance"] == 3

From fb35a1d0328cd19858da623bd5d00f08d9f8a09e Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 12 Jan 2026 00:47:23 +0800
Subject: [PATCH 02/33] Add ImageNet LPIPS benchmark workflow

---
 .github/workflows/benchmark.yml |  66 +++++++++
 tests/benchmark_lpips.py        | 228 ++++++++++++++++++++++++++++++++
 tests/compare_benchmarks.py     |  63 +++++++++
 3 files changed, 357 insertions(+)
 create mode 100644 .github/workflows/benchmark.yml
 create mode 100644 tests/benchmark_lpips.py
 create mode 100644 tests/compare_benchmarks.py

diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml
new file mode 100644
index 0000000..9d517a8
--- /dev/null
+++ b/.github/workflows/benchmark.yml
@@ -0,0 +1,66 @@
+name: LanPaintBench LPIPS (ImageNet subset)
+
+on:
+  pull_request:
+    branches: [ master, main ]
+    paths:
+      - 'src/**'
+      - 'tests/benchmark_lpips.py'
+      - 'tests/compare_benchmarks.py'
+      - '.github/workflows/benchmark.yml'
+  workflow_dispatch:
+
+jobs:
+  benchmark:
+    name: Compare Base vs Head
+    runs-on: ubuntu-latest
+    env:
+      HF_HOME: ~/.cache/huggingface
+      BASE_REF: ${{ github.event.pull_request.base.sha || github.event.repository.default_branch }}
+    steps:
+      - name: Checkout Head (Current PR)
+        uses: actions/checkout@v4
+        with:
+          path: head
+          fetch-depth: 1
+
+      - name: Checkout Base (Baseline)
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ env.BASE_REF }}
+          path: base
+          fetch-depth: 1
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+          cache: 'pip'
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+          pip install lpips datasets numpy pillow
+          # We don't need to install LanPaint as a package because we use sys.path in the script
+
+      - name: Cache HF + torch
+        uses: actions/cache@v4
+        with:
+          path: |
+            ~/.cache/huggingface
+            ~/.cache/torch
+          key: ${{ runner.os }}-lpips-imagenet-${{ hashFiles('head/tests/benchmark_lpips.py', 'head/tests/compare_benchmarks.py') }}
+
+      - name: Run Baseline Benchmark
+        run: |
+          # Use the benchmark script from the HEAD branch to ensure test consistency
+          python head/tests/benchmark_lpips.py --src base/src --output baseline.json --subset 8 --seed 42 --mask-type box --n-steps 5 --sigma 1.0 --device cpu
+
+      - name: Run Current Benchmark
+        run: |
+          python head/tests/benchmark_lpips.py --src head/src --output current.json --subset 8 --seed 42 --mask-type box --n-steps 5 --sigma 1.0 --device cpu
+
+      - name: Compare Results
+        run: |
+          python head/tests/compare_benchmarks.py --baseline baseline.json --current current.json --max-rel-diff 0.01 --max-abs-diff 0.0
diff --git a/tests/benchmark_lpips.py b/tests/benchmark_lpips.py
new file mode 100644
index 0000000..f3a7aa8
--- /dev/null
+++ b/tests/benchmark_lpips.py
@@ -0,0 +1,228 @@
+"""Lightweight CI benchmark inspired by LanPaintBench.
+
+Runs LPIPS on a tiny deterministic ImageNet-1k-256x256 subset and compares base vs head.
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+def _make_mask(mask_type: str, height: int, width: int) -> torch.Tensor:
+    """Return a (1, 3, H, W) mask where 1=known region, 0=to-inpaint."""
+    if height <= 0 or width <= 0:
+        raise ValueError("height/width must be positive")
+
+    mask = torch.ones((1, 3, height, width), dtype=torch.float32)
+
+    if mask_type == "half":
+        mask[:, :, :, width // 2 :] = 0.0
+        return mask
+
+    if mask_type == "box":
+        h0, h1 = height // 4, (3 * height) // 4
+        w0, w1 = width // 4, (3 * width) // 4
+        mask[:, :, h0:h1, w0:w1] = 0.0
+        return mask
+
+    if mask_type == "outpaint":
+        mask = torch.zeros((1, 3, height, width), dtype=torch.float32)
+        h0, h1 = height // 4, (3 * height) // 4
+        w0, w1 = width // 4, (3 * width) // 4
+        mask[:, :, h0:h1, w0:w1] = 1.0
+        return mask
+
+    if mask_type == "checkerboard":
+        grid = max(1, min(height, width) // 16)
+        mask_2d = torch.zeros((height, width), dtype=torch.float32)
+        for row in range(height):
+            row_parity = (row // grid) % 2
+            start = 0 if row_parity == 0 else 1
+            for block in range(start, (width + grid - 1) // grid, 2):
+                c0 = block * grid
+                c1 = min(width, (block + 1) * grid)
+                mask_2d[row, c0:c1] = 1.0
+        return mask_2d.unsqueeze(0).unsqueeze(0).repeat(1, 3, 1, 1)
+
+    raise ValueError(f"unknown mask type: {mask_type}")
+
+
+class _DummySampling:
+    def noise_scaling(self, sigma, noise, latent_image):  # type: ignore[no-untyped-def]
+        # Variance exploding convention (ComfyUI/k-diffusion style).
+        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]
+        # Cheap "denoiser" so the benchmark exercises LanPaint math deterministically on CPU.
+        sigma_ = sigma.reshape((sigma.shape[0],) + (1,) * (x.ndim - 1))
+        x0 = x / (1.0 + sigma_**2)
+        return x0, x0
+
+
+def _load_images(*, dataset: str, split: str, subset: int, seed: int, shuffle_buffer: int) -> list[torch.Tensor]:
+    from datasets import load_dataset
+
+    ds = load_dataset(dataset, split=split, streaming=True)
+    try:
+        ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer)
+    except Exception:
+        pass
+
+    images: list[torch.Tensor] = []
+    for idx, item in enumerate(ds):
+        if idx >= subset:
+            break
+        image = item["image"].convert("RGB")
+        arr = np.asarray(image, dtype=np.uint8)
+        img = torch.from_numpy(arr).permute(2, 0, 1).float() / 127.5 - 1.0
+        images.append(img)
+    return images
+
+
+def _run_bench(  # type: ignore[no-untyped-def]
+    LanPaintClass,
+    images: list[torch.Tensor],
+    device: torch.device,
+    *,
+    n_steps: int,
+    sigma_value: float,
+    mask_type: str,
+    seed: int,
+) -> float:
+    import lpips
+
+    lpips_fn = lpips.LPIPS(net="alex").to(device)
+
+    sigma = torch.tensor([sigma_value], device=device)
+    abt = 1.0 / (1.0 + sigma**2)
+    flow_t = torch.zeros_like(sigma)
+    current_times = (sigma, abt, flow_t)
+
+    engine = LanPaintClass(
+        Model=_DummyModel(),
+        NSteps=n_steps,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.1,
+    )
+
+    scores: list[float] = []
+    for idx, img in enumerate(images):
+        img = img.unsqueeze(0).to(device)
+        latent_image = img.clone()
+
+        height, width = img.shape[-2:]
+        latent_mask = _make_mask(mask_type, height, width).to(device)
+
+        gen = torch.Generator(device=device.type).manual_seed(seed + idx)
+        noise = torch.randn(img.shape, generator=gen, device=device)
+        x = noise * sigma.reshape((1,) + (1,) * (img.ndim - 1))
+
+        output = engine(
+            x=x,
+            latent_image=latent_image,
+            noise=noise,
+            sigma=sigma,
+            latent_mask=latent_mask,
+            current_times=current_times,
+            model_options={},
+            seed=seed,
+            n_steps=n_steps,
+        ).clamp(-1.0, 1.0)
+
+        dist = lpips_fn(output, img)
+        scores.append(float(dist.item()))
+
+    return float(np.mean(scores))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="LanPaint LPIPS benchmark (ImageNet-1k-256x256 subset)")
+    parser.add_argument("--src", type=Path, required=True, help="Path to LanPaint src directory (contains LanPaint/)")
+    parser.add_argument("--output", type=Path, required=True, help="Output JSON path")
+
+    parser.add_argument("--dataset", type=str, default="benjamin-paine/imagenet-1k-256x256")
+    parser.add_argument("--split", type=str, default="train")
+    parser.add_argument("--subset", type=int, default=8)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--shuffle-buffer", type=int, default=1024)
+
+    parser.add_argument("--mask-type", type=str, default="box", choices=["half", "box", "outpaint", "checkerboard"])
+    parser.add_argument("--n-steps", type=int, default=5, help="LanPaint inner iterations per call")
+    parser.add_argument("--sigma", type=float, default=1.0, help="Variance-exploding sigma value")
+    parser.add_argument("--device", type=str, default="cpu", choices=["cpu", "cuda", "auto"])
+    args = parser.parse_args()
+
+    if args.subset <= 0:
+        raise SystemExit("--subset must be > 0")
+    if args.n_steps <= 0:
+        raise SystemExit("--n-steps must be > 0")
+    if args.shuffle_buffer <= 0:
+        raise SystemExit("--shuffle-buffer must be > 0")
+
+    device_str = args.device
+    if device_str == "auto":
+        device_str = "cuda" if torch.cuda.is_available() else "cpu"
+    device = torch.device(device_str)
+
+    sys.path.insert(0, str(args.src.resolve()))
+    from LanPaint.lanpaint import LanPaint as LanPaintEngine
+
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    images = _load_images(
+        dataset=args.dataset,
+        split=args.split,
+        subset=args.subset,
+        seed=args.seed,
+        shuffle_buffer=args.shuffle_buffer,
+    )
+
+    score = _run_bench(
+        LanPaintEngine,
+        images,
+        device,
+        n_steps=args.n_steps,
+        sigma_value=args.sigma,
+        mask_type=args.mask_type,
+        seed=args.seed,
+    )
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with args.output.open("w", encoding="utf-8") as f:
+        json.dump(
+            {
+                "lpips": score,
+                "config": {
+                    "dataset": args.dataset,
+                    "split": args.split,
+                    "subset": args.subset,
+                    "seed": args.seed,
+                    "shuffle_buffer": args.shuffle_buffer,
+                    "mask_type": args.mask_type,
+                    "n_steps": args.n_steps,
+                    "sigma": args.sigma,
+                    "device": device_str,
+                },
+            },
+            f,
+            indent=2,
+        )
+
+    print(f"LPIPS: {score:.6f}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/compare_benchmarks.py b/tests/compare_benchmarks.py
new file mode 100644
index 0000000..8846143
--- /dev/null
+++ b/tests/compare_benchmarks.py
@@ -0,0 +1,63 @@
+"""Compare two LPIPS benchmark JSON outputs (baseline vs current)."""
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+
+def _append_step_summary(markdown: str) -> None:
+    summary_path = os.environ.get("GITHUB_STEP_SUMMARY")
+    if not summary_path:
+        return
+    with Path(summary_path).open("a", encoding="utf-8", errors="ignore") as f:
+        f.write(markdown)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--baseline", type=Path, required=True)
+    parser.add_argument("--current", type=Path, required=True)
+    parser.add_argument("--max-rel-diff", type=float, default=0.01, help="Fail if (current-baseline)/baseline exceeds this")
+    parser.add_argument("--max-abs-diff", type=float, default=0.0, help="Fail if current-baseline exceeds this")
+    args = parser.parse_args()
+
+    baseline = json.loads(args.baseline.read_text(encoding="utf-8"))
+    current = json.loads(args.current.read_text(encoding="utf-8"))
+
+    b_score = float(baseline["lpips"])
+    c_score = float(current["lpips"])
+    diff = c_score - b_score
+    rel_diff = diff / b_score if b_score != 0.0 else float("inf")
+
+    lines = [
+        f"Baseline LPIPS: {b_score:.6f}",
+        f"Current LPIPS:  {c_score:.6f}",
+        f"Difference:     {diff:+.6f}",
+    ]
+    if b_score != 0.0:
+        lines.append(f"Relative diff:   {rel_diff:+.2%}")
+    print("\n".join(lines))
+
+    summary = "\n".join(
+        [
+            "### LPIPS benchmark",
+            "",
+            "| Metric | Baseline | Current | Diff |",
+            "|---|---:|---:|---:|",
+            f"| LPIPS | {b_score:.6f} | {c_score:.6f} | {diff:+.6f} |",
+            "",
+        ]
+    )
+    _append_step_summary(summary)
+
+    if diff > args.max_abs_diff:
+        if b_score == 0.0 or rel_diff > args.max_rel_diff:
+            print("FAIL: LPIPS regression exceeds configured threshold(s)")
+            raise SystemExit(1)
+
+    print("PASS: LPIPS within configured threshold(s)")
+
+
+if __name__ == "__main__":
+    main()

From 7c79712fc834c276bff4fbafb801bcf1ca1bd01f Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 12 Jan 2026 09:17:29 +0800
Subject: [PATCH 03/33] Make LPIPS benchmark long-running e2e

---
 .github/workflows/benchmark.yml  |  68 ++++++--
 tests/benchmark_lpips.py         | 276 +++++++++++++++++++++++++++++--
 tests/compare_benchmarks.py      |  10 ++
 tests/test_benchmark_schedule.py |  25 +++
 4 files changed, 354 insertions(+), 25 deletions(-)
 create mode 100644 tests/test_benchmark_schedule.py

diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml
index 9d517a8..bcadcbd 100644
--- a/.github/workflows/benchmark.yml
+++ b/.github/workflows/benchmark.yml
@@ -1,22 +1,20 @@
 name: LanPaintBench LPIPS (ImageNet subset)
 
 on:
-  pull_request:
-    branches: [ master, main ]
-    paths:
-      - 'src/**'
-      - 'tests/benchmark_lpips.py'
-      - 'tests/compare_benchmarks.py'
-      - '.github/workflows/benchmark.yml'
   workflow_dispatch:
 
 jobs:
   benchmark:
     name: Compare Base vs Head
     runs-on: ubuntu-latest
+    timeout-minutes: 360
     env:
       HF_HOME: ~/.cache/huggingface
       BASE_REF: ${{ github.event.pull_request.base.sha || github.event.repository.default_branch }}
+      OMP_NUM_THREADS: "1"
+      MKL_NUM_THREADS: "1"
+      LANPAINTBENCH_REF: "ea5be5d"
+      CHECKPOINT_URL: "https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion_uncond.pt"
     steps:
       - name: Checkout Head (Current PR)
         uses: actions/checkout@v4
@@ -41,25 +39,73 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
-          pip install lpips datasets numpy pillow
+          pip install lpips datasets numpy pillow blobfile
           # We don't need to install LanPaint as a package because we use sys.path in the script
 
-      - name: Cache HF + torch
+      - name: Cache HF + torch + checkpoints
         uses: actions/cache@v4
         with:
           path: |
             ~/.cache/huggingface
             ~/.cache/torch
+            ~/.cache/lanpaintbench
           key: ${{ runner.os }}-lpips-imagenet-${{ hashFiles('head/tests/benchmark_lpips.py', 'head/tests/compare_benchmarks.py') }}
 
+      - name: Clone LanPaintBench
+        run: |
+          git clone https://github.com/scraed/LanPaintBench.git bench
+          cd bench
+          git checkout "$LANPAINTBENCH_REF"
+
+      - name: Download ImageNet-256 checkpoint
+        run: |
+          CHECKPOINT_PATH="$HOME/.cache/lanpaintbench/checkpoints/256x256_diffusion_uncond.pt"
+          mkdir -p ~/.cache/lanpaintbench/checkpoints
+          if [ ! -f "$CHECKPOINT_PATH" ]; then
+            curl -L "$CHECKPOINT_URL" -o "$CHECKPOINT_PATH"
+          fi
+          ls -lh "$CHECKPOINT_PATH"
+
       - name: Run Baseline Benchmark
         run: |
+          CHECKPOINT_PATH="$HOME/.cache/lanpaintbench/checkpoints/256x256_diffusion_uncond.pt"
           # Use the benchmark script from the HEAD branch to ensure test consistency
-          python head/tests/benchmark_lpips.py --src base/src --output baseline.json --subset 8 --seed 42 --mask-type box --n-steps 5 --sigma 1.0 --device cpu
+          python head/tests/benchmark_lpips.py \
+            --mode guided_diffusion_e2e \
+            --src base/src \
+            --lanpaintbench bench \
+            --checkpoint "$CHECKPOINT_PATH" \
+            --output baseline.json \
+            --dataset benjamin-paine/imagenet-1k-256x256 \
+            --split train \
+            --subset 50 \
+            --seed 42 \
+            --shuffle-buffer 1024 \
+            --mask-type box \
+            --n-steps 5 \
+            --sample-steps 20 \
+            --device cpu \
+            --min-seconds 3600
 
       - name: Run Current Benchmark
         run: |
-          python head/tests/benchmark_lpips.py --src head/src --output current.json --subset 8 --seed 42 --mask-type box --n-steps 5 --sigma 1.0 --device cpu
+          CHECKPOINT_PATH="$HOME/.cache/lanpaintbench/checkpoints/256x256_diffusion_uncond.pt"
+          python head/tests/benchmark_lpips.py \
+            --mode guided_diffusion_e2e \
+            --src head/src \
+            --lanpaintbench bench \
+            --checkpoint "$CHECKPOINT_PATH" \
+            --output current.json \
+            --dataset benjamin-paine/imagenet-1k-256x256 \
+            --split train \
+            --subset 50 \
+            --seed 42 \
+            --shuffle-buffer 1024 \
+            --mask-type box \
+            --n-steps 5 \
+            --sample-steps 20 \
+            --device cpu \
+            --min-seconds 3600
 
       - name: Compare Results
         run: |
diff --git a/tests/benchmark_lpips.py b/tests/benchmark_lpips.py
index f3a7aa8..d85c499 100644
--- a/tests/benchmark_lpips.py
+++ b/tests/benchmark_lpips.py
@@ -1,11 +1,14 @@
-"""Lightweight CI benchmark inspired by LanPaintBench.
+"""LanPaintBench-style LPIPS benchmark for GitHub Actions.
 
-Runs LPIPS on a tiny deterministic ImageNet-1k-256x256 subset and compares base vs head.
+Supports:
+- A fast dummy-mode sanity benchmark.
+- A slow e2e-ish mode using LanPaintBench's `guided_diffusion` ImageNet-256 checkpoint.
 """
 
 import argparse
 import json
 import sys
+import time
 from pathlib import Path
 
 import numpy as np
@@ -89,7 +92,7 @@ def _load_images(*, dataset: str, split: str, subset: int, seed: int, shuffle_bu
     return images
 
 
-def _run_bench(  # type: ignore[no-untyped-def]
+def _run_dummy_bench(  # type: ignore[no-untyped-def]
     LanPaintClass,
     images: list[torch.Tensor],
     device: torch.device,
@@ -147,10 +150,217 @@ def _run_bench(  # type: ignore[no-untyped-def]
     return float(np.mean(scores))
 
 
+def _ddpm_alpha_bar(*, timesteps: int, beta_start: float, beta_end: float) -> torch.Tensor:
+    betas = torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float64)
+    alphas = 1.0 - betas
+    return torch.cumprod(alphas, dim=0)
+
+
+def _sigmas_and_timesteps(
+    *,
+    inference_steps: int,
+    train_timesteps: int,
+    beta_start: float,
+    beta_end: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if inference_steps <= 0:
+        raise ValueError("inference_steps must be > 0")
+    if train_timesteps <= 1:
+        raise ValueError("train_timesteps must be > 1")
+
+    alpha_bar = _ddpm_alpha_bar(timesteps=train_timesteps, beta_start=beta_start, beta_end=beta_end)
+    t = torch.linspace(train_timesteps - 1, 0, inference_steps, dtype=torch.float64).round().to(torch.long)
+    t = torch.unique_consecutive(t)
+    if t.numel() < 2:
+        raise ValueError("inference_steps too small after rounding")
+
+    abt = alpha_bar[t].clamp(min=1e-12, max=1.0)
+    sigmas = torch.sqrt((1.0 - abt) / abt).to(torch.float32)
+    sigmas = torch.cat([sigmas, sigmas.new_zeros(1)], dim=0)
+    return sigmas, t.to(torch.long)
+
+
+class _GuidedDiffusionSampling:
+    def noise_scaling(self, sigma, noise, latent_image):  # type: ignore[no-untyped-def]
+        return latent_image + noise * sigma
+
+
+class _GuidedDiffusionModel:
+    def __init__(self, *, unet, device: torch.device) -> None:  # type: ignore[no-untyped-def]
+        self.inner_model = self
+        self.model_sampling = _GuidedDiffusionSampling()
+        self._unet = unet
+        self._device = device
+
+    @torch.no_grad()
+    def __call__(self, x, sigma, model_options=None, seed=None):  # type: ignore[no-untyped-def]
+        if not isinstance(model_options, dict) or "bench_timestep" not in model_options:
+            raise ValueError("model_options['bench_timestep'] is required for guided_diffusion mode")
+
+        t_int = int(model_options["bench_timestep"])
+        t = torch.full((x.shape[0],), t_int, device=self._device, dtype=torch.long)
+
+        sigma_ = sigma.reshape((sigma.shape[0],) + (1,) * (x.ndim - 1))
+        abt = 1.0 / (1.0 + sigma_**2)
+
+        x_vp = x / torch.sqrt(1.0 + sigma_**2)
+        eps = self._unet(x_vp, t)
+        eps = eps[:, :3, :, :]
+
+        sqrt_abt = torch.sqrt(abt)
+        sqrt_1m_abt = torch.sqrt(1.0 - abt)
+        x0 = (x_vp - sqrt_1m_abt * eps) / sqrt_abt
+        return x0, x0
+
+
+def _load_guided_diffusion_unet(*, lanpaintbench: Path, checkpoint: Path, device: torch.device):  # type: ignore[no-untyped-def]
+    sys.path.insert(0, str(lanpaintbench.resolve()))
+    from guided_diffusion.unet import UNetModel  # type: ignore[import-not-found]
+
+    image_size = 256
+    num_channels = 256
+    num_head_channels = 64
+    num_res_blocks = 2
+    resblock_updown = True
+    learn_sigma = True
+    class_cond = False
+    use_checkpoint = False
+    attention_resolutions = "32,16,8"
+    use_scale_shift_norm = True
+    use_fp16 = False
+
+    channel_mult = (1, 1, 2, 2, 4, 4)
+    attention_ds = []
+    for res in attention_resolutions.split(","):
+        attention_ds.append(image_size // int(res))
+
+    model = UNetModel(
+        image_size=image_size,
+        in_channels=3,
+        model_channels=num_channels,
+        out_channels=(3 if not learn_sigma else 6),
+        num_res_blocks=num_res_blocks,
+        attention_resolutions=tuple(attention_ds),
+        dropout=0.0,
+        channel_mult=channel_mult,
+        num_classes=(1000 if class_cond else None),
+        use_checkpoint=use_checkpoint,
+        use_fp16=use_fp16,
+        num_heads=1,
+        num_head_channels=num_head_channels,
+        num_heads_upsample=-1,
+        use_scale_shift_norm=use_scale_shift_norm,
+        resblock_updown=resblock_updown,
+        use_new_attention_order=False,
+    ).to(device)
+
+    state = torch.load(checkpoint, map_location=device)
+    model.load_state_dict(state)
+    model.eval()
+    return model
+
+
+def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
+    LanPaintClass,
+    images: list[torch.Tensor],
+    device: torch.device,
+    *,
+    lanpaintbench: Path,
+    checkpoint: Path,
+    seed: int,
+    mask_type: str,
+    sample_steps: int,
+    inner_steps: int,
+    beta_start: float,
+    beta_end: float,
+    train_timesteps: int,
+) -> tuple[float, dict[str, float]]:
+    import lpips
+
+    unet = _load_guided_diffusion_unet(lanpaintbench=lanpaintbench, checkpoint=checkpoint, device=device)
+    model = _GuidedDiffusionModel(unet=unet, device=device)
+
+    sigmas, timesteps = _sigmas_and_timesteps(
+        inference_steps=sample_steps,
+        train_timesteps=train_timesteps,
+        beta_start=beta_start,
+        beta_end=beta_end,
+    )
+    sigmas = sigmas.to(device=device)
+    timesteps = timesteps.to(device=device)
+
+    lpips_fn = lpips.LPIPS(net="alex").to(device)
+
+    scores: list[float] = []
+    start = time.perf_counter()
+
+    engine = LanPaintClass(
+        Model=model,
+        NSteps=inner_steps,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.1,
+    )
+
+    for idx, img in enumerate(images):
+        img = img.unsqueeze(0).to(device)
+
+        height, width = img.shape[-2:]
+        latent_mask = _make_mask(mask_type, height, width).to(device)
+        inpaint_mask = 1.0 - latent_mask
+
+        latent_image = img * latent_mask
+
+        gen = torch.Generator(device=device.type).manual_seed(seed + idx)
+        noise = torch.randn(img.shape, generator=gen, device=device)
+
+        x = latent_image + noise * sigmas[0].reshape((1,) + (1,) * (img.ndim - 1))
+
+        for step_idx in range(sigmas.shape[0] - 1):
+            sigma = sigmas[step_idx : step_idx + 1]
+            sigma_next = sigmas[step_idx + 1 : step_idx + 2]
+            t = timesteps[min(step_idx, timesteps.shape[0] - 1)].item()
+
+            abt = 1.0 / (1.0 + sigma**2)
+            flow_t = torch.sqrt(1.0 - abt) / (torch.sqrt(1.0 - abt) + torch.sqrt(abt))
+            current_times = (sigma, abt, flow_t)
+
+            x0 = engine(
+                x=x,
+                latent_image=latent_image,
+                noise=noise,
+                sigma=sigma,
+                latent_mask=latent_mask,
+                current_times=current_times,
+                model_options={"bench_timestep": int(t)},
+                seed=seed,
+                n_steps=inner_steps,
+            )
+
+            if float(sigma_next.item()) == 0.0:
+                x = x0
+                break
+
+            sigma_ = sigma.reshape((sigma.shape[0],) + (1,) * (x.ndim - 1))
+            sigma_next_ = sigma_next.reshape((sigma_next.shape[0],) + (1,) * (x.ndim - 1))
+            d = (x - x0) / sigma_
+            x = x + d * (sigma_next_ - sigma_)
+
+        x = x.clamp(-1.0, 1.0)
+        dist = lpips_fn(x * inpaint_mask, img * inpaint_mask)
+        scores.append(float(dist.item()))
+
+    elapsed = time.perf_counter() - start
+    stats = {"seconds_total": elapsed, "seconds_per_image": elapsed / max(1, len(images))}
+    return float(np.mean(scores)), stats
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(description="LanPaint LPIPS benchmark (ImageNet-1k-256x256 subset)")
     parser.add_argument("--src", type=Path, required=True, help="Path to LanPaint src directory (contains LanPaint/)")
     parser.add_argument("--output", type=Path, required=True, help="Output JSON path")
+    parser.add_argument("--mode", type=str, default="dummy", choices=["dummy", "guided_diffusion_e2e"])
 
     parser.add_argument("--dataset", type=str, default="benjamin-paine/imagenet-1k-256x256")
     parser.add_argument("--split", type=str, default="train")
@@ -159,9 +369,17 @@ def main() -> None:
     parser.add_argument("--shuffle-buffer", type=int, default=1024)
 
     parser.add_argument("--mask-type", type=str, default="box", choices=["half", "box", "outpaint", "checkerboard"])
-    parser.add_argument("--n-steps", type=int, default=5, help="LanPaint inner iterations per call")
-    parser.add_argument("--sigma", type=float, default=1.0, help="Variance-exploding sigma value")
+    parser.add_argument("--n-steps", type=int, default=5, help="LanPaint inner iterations per call (dummy mode)")
+    parser.add_argument("--sigma", type=float, default=1.0, help="Variance-exploding sigma value (dummy mode)")
     parser.add_argument("--device", type=str, default="cpu", choices=["cpu", "cuda", "auto"])
+
+    parser.add_argument("--lanpaintbench", type=Path, default=None, help="Path to LanPaintBench checkout (guided_diffusion_e2e)")
+    parser.add_argument("--checkpoint", type=Path, default=None, help="Path to 256x256_diffusion_uncond.pt (guided_diffusion_e2e)")
+    parser.add_argument("--sample-steps", type=int, default=20, help="Outer diffusion steps (guided_diffusion_e2e)")
+    parser.add_argument("--train-timesteps", type=int, default=1000, help="Training timesteps for beta schedule (guided_diffusion_e2e)")
+    parser.add_argument("--beta-start", type=float, default=0.0001, help="DDPM beta_start (guided_diffusion_e2e)")
+    parser.add_argument("--beta-end", type=float, default=0.02, help="DDPM beta_end (guided_diffusion_e2e)")
+    parser.add_argument("--min-seconds", type=float, default=0.0, help="Fail if benchmark runtime is below this (guided_diffusion_e2e)")
     args = parser.parse_args()
 
     if args.subset <= 0:
@@ -190,22 +408,46 @@ def main() -> None:
         shuffle_buffer=args.shuffle_buffer,
     )
 
-    score = _run_bench(
-        LanPaintEngine,
-        images,
-        device,
-        n_steps=args.n_steps,
-        sigma_value=args.sigma,
-        mask_type=args.mask_type,
-        seed=args.seed,
-    )
+    stats: dict[str, float] = {}
+    if args.mode == "dummy":
+        score = _run_dummy_bench(
+            LanPaintEngine,
+            images,
+            device,
+            n_steps=args.n_steps,
+            sigma_value=args.sigma,
+            mask_type=args.mask_type,
+            seed=args.seed,
+        )
+    else:
+        if args.lanpaintbench is None or args.checkpoint is None:
+            raise SystemExit("--lanpaintbench and --checkpoint are required for guided_diffusion_e2e mode")
+
+        score, stats = _run_guided_diffusion_e2e_bench(
+            LanPaintEngine,
+            images,
+            device,
+            lanpaintbench=args.lanpaintbench,
+            checkpoint=args.checkpoint,
+            seed=args.seed,
+            mask_type=args.mask_type,
+            sample_steps=args.sample_steps,
+            inner_steps=args.n_steps,
+            beta_start=args.beta_start,
+            beta_end=args.beta_end,
+            train_timesteps=args.train_timesteps,
+        )
+        if args.min_seconds > 0.0 and stats.get("seconds_total", 0.0) < args.min_seconds:
+            raise SystemExit(f"benchmark too fast: {stats.get('seconds_total', 0.0):.1f}s < {args.min_seconds:.1f}s")
 
     args.output.parent.mkdir(parents=True, exist_ok=True)
     with args.output.open("w", encoding="utf-8") as f:
         json.dump(
             {
                 "lpips": score,
+                "stats": stats,
                 "config": {
+                    "mode": args.mode,
                     "dataset": args.dataset,
                     "split": args.split,
                     "subset": args.subset,
@@ -214,6 +456,12 @@ def main() -> None:
                     "mask_type": args.mask_type,
                     "n_steps": args.n_steps,
                     "sigma": args.sigma,
+                    "sample_steps": args.sample_steps,
+                    "train_timesteps": args.train_timesteps,
+                    "beta_start": args.beta_start,
+                    "beta_end": args.beta_end,
+                    "lanpaintbench": str(args.lanpaintbench) if args.lanpaintbench is not None else None,
+                    "checkpoint": str(args.checkpoint) if args.checkpoint is not None else None,
                     "device": device_str,
                 },
             },
diff --git a/tests/compare_benchmarks.py b/tests/compare_benchmarks.py
index 8846143..f90d595 100644
--- a/tests/compare_benchmarks.py
+++ b/tests/compare_benchmarks.py
@@ -27,6 +27,14 @@ def main() -> None:
 
     b_score = float(baseline["lpips"])
     c_score = float(current["lpips"])
+
+    b_stats = baseline.get("stats") if isinstance(baseline, dict) else None
+    c_stats = current.get("stats") if isinstance(current, dict) else None
+    b_seconds = float(b_stats.get("seconds_total", 0.0)) if isinstance(b_stats, dict) else 0.0
+    c_seconds = float(c_stats.get("seconds_total", 0.0)) if isinstance(c_stats, dict) else 0.0
+    b_spi = float(b_stats.get("seconds_per_image", 0.0)) if isinstance(b_stats, dict) else 0.0
+    c_spi = float(c_stats.get("seconds_per_image", 0.0)) if isinstance(c_stats, dict) else 0.0
+
     diff = c_score - b_score
     rel_diff = diff / b_score if b_score != 0.0 else float("inf")
 
@@ -46,6 +54,8 @@ def main() -> None:
             "| Metric | Baseline | Current | Diff |",
             "|---|---:|---:|---:|",
             f"| LPIPS | {b_score:.6f} | {c_score:.6f} | {diff:+.6f} |",
+            f"| Runtime (s) | {b_seconds:.1f} | {c_seconds:.1f} | {(c_seconds - b_seconds):+.1f} |",
+            f"| Sec/img | {b_spi:.2f} | {c_spi:.2f} | {(c_spi - b_spi):+.2f} |",
             "",
         ]
     )
diff --git a/tests/test_benchmark_schedule.py b/tests/test_benchmark_schedule.py
new file mode 100644
index 0000000..661f6cc
--- /dev/null
+++ b/tests/test_benchmark_schedule.py
@@ -0,0 +1,25 @@
+import torch
+
+from tests.benchmark_lpips import _make_mask, _sigmas_and_timesteps
+
+
+def test_sigmas_schedule_descends_and_ends_at_zero() -> None:
+    sigmas, timesteps = _sigmas_and_timesteps(
+        inference_steps=20,
+        train_timesteps=1000,
+        beta_start=0.0001,
+        beta_end=0.02,
+    )
+
+    assert sigmas.numel() == timesteps.numel() + 1
+    assert sigmas[-1].item() == 0.0
+    assert torch.all(sigmas[:-1] > 0.0)
+    assert torch.all(sigmas[:-1] >= sigmas[1:])
+    assert timesteps[0].item() >= timesteps[-1].item()
+
+
+def test_make_mask_shapes() -> None:
+    mask = _make_mask("box", 64, 64)
+    assert tuple(mask.shape) == (1, 3, 64, 64)
+    assert mask.dtype == torch.float32
+

From 11c1490140b9f9b6661969361a4f8796cc6978d3 Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Wed, 14 Jan 2026 11:26:54 +0800
Subject: [PATCH 04/33] Make benchmark workflow test semantic early stop

---
 .github/workflows/benchmark.yml | 121 +++++++++++++++++++++++++++-----
 tests/benchmark_lpips.py        |  31 +++++++-
 2 files changed, 134 insertions(+), 18 deletions(-)

diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml
index bcadcbd..e753772 100644
--- a/.github/workflows/benchmark.yml
+++ b/.github/workflows/benchmark.yml
@@ -2,6 +2,47 @@ name: LanPaintBench LPIPS (ImageNet subset)
 
 on:
   workflow_dispatch:
+    inputs:
+      base_ref:
+        description: "Base ref to benchmark against (commit SHA or branch). Leave empty to use repo default branch."
+        required: false
+        default: ""
+      subset:
+        description: "Number of ImageNet images to benchmark"
+        required: true
+        default: "10"
+      sample_steps:
+        description: "Outer diffusion steps (guided_diffusion_e2e)"
+        required: true
+        default: "20"
+      n_steps:
+        description: "LanPaint inner steps per outer step"
+        required: true
+        default: "5"
+      seed:
+        description: "Random seed"
+        required: true
+        default: "42"
+      shuffle_buffer:
+        description: "HF streaming shuffle buffer"
+        required: true
+        default: "1024"
+      semantic_threshold:
+        description: "Enable semantic early-stop if > 0 (head_on run)"
+        required: true
+        default: "0.01"
+      semantic_patience:
+        description: "Semantic early-stop patience (head_on run)"
+        required: true
+        default: "2"
+      semantic_min_steps:
+        description: "Semantic early-stop minimum steps (head_on run)"
+        required: true
+        default: "2"
+      min_seconds:
+        description: "Fail if a benchmark runtime is below this (0 to disable)"
+        required: true
+        default: "0"
 
 jobs:
   benchmark:
@@ -10,7 +51,7 @@ jobs:
     timeout-minutes: 360
     env:
       HF_HOME: ~/.cache/huggingface
-      BASE_REF: ${{ github.event.pull_request.base.sha || github.event.repository.default_branch }}
+      BASE_REF: ${{ inputs.base_ref || github.event.repository.default_branch }}
       OMP_NUM_THREADS: "1"
       MKL_NUM_THREADS: "1"
       LANPAINTBENCH_REF: "ea5be5d"
@@ -78,16 +119,17 @@ jobs:
             --output baseline.json \
             --dataset benjamin-paine/imagenet-1k-256x256 \
             --split train \
-            --subset 50 \
-            --seed 42 \
-            --shuffle-buffer 1024 \
+            --subset "${{ inputs.subset }}" \
+            --seed "${{ inputs.seed }}" \
+            --shuffle-buffer "${{ inputs.shuffle_buffer }}" \
             --mask-type box \
-            --n-steps 5 \
-            --sample-steps 20 \
+            --n-steps "${{ inputs.n_steps }}" \
+            --sample-steps "${{ inputs.sample_steps }}" \
             --device cpu \
-            --min-seconds 3600
+            --min-seconds "${{ inputs.min_seconds }}" \
+            --semantic-threshold 0.0
 
-      - name: Run Current Benchmark
+      - name: Run Head Benchmark (early-stop OFF)
         run: |
           CHECKPOINT_PATH="$HOME/.cache/lanpaintbench/checkpoints/256x256_diffusion_uncond.pt"
           python head/tests/benchmark_lpips.py \
@@ -95,18 +137,63 @@ jobs:
             --src head/src \
             --lanpaintbench bench \
             --checkpoint "$CHECKPOINT_PATH" \
-            --output current.json \
+            --output head_off.json \
             --dataset benjamin-paine/imagenet-1k-256x256 \
             --split train \
-            --subset 50 \
-            --seed 42 \
-            --shuffle-buffer 1024 \
+            --subset "${{ inputs.subset }}" \
+            --seed "${{ inputs.seed }}" \
+            --shuffle-buffer "${{ inputs.shuffle_buffer }}" \
             --mask-type box \
-            --n-steps 5 \
-            --sample-steps 20 \
+            --n-steps "${{ inputs.n_steps }}" \
+            --sample-steps "${{ inputs.sample_steps }}" \
             --device cpu \
-            --min-seconds 3600
+            --min-seconds "${{ inputs.min_seconds }}" \
+            --semantic-threshold 0.0
 
-      - name: Compare Results
+      - name: Run Head Benchmark (early-stop ON)
         run: |
-          python head/tests/compare_benchmarks.py --baseline baseline.json --current current.json --max-rel-diff 0.01 --max-abs-diff 0.0
+          CHECKPOINT_PATH="$HOME/.cache/lanpaintbench/checkpoints/256x256_diffusion_uncond.pt"
+          python head/tests/benchmark_lpips.py \
+            --mode guided_diffusion_e2e \
+            --src head/src \
+            --lanpaintbench bench \
+            --checkpoint "$CHECKPOINT_PATH" \
+            --output head_on.json \
+            --dataset benjamin-paine/imagenet-1k-256x256 \
+            --split train \
+            --subset "${{ inputs.subset }}" \
+            --seed "${{ inputs.seed }}" \
+            --shuffle-buffer "${{ inputs.shuffle_buffer }}" \
+            --mask-type box \
+            --n-steps "${{ inputs.n_steps }}" \
+            --sample-steps "${{ inputs.sample_steps }}" \
+            --device cpu \
+            --min-seconds "${{ inputs.min_seconds }}" \
+            --semantic-threshold "${{ inputs.semantic_threshold }}" \
+            --semantic-patience "${{ inputs.semantic_patience }}" \
+            --semantic-min-steps "${{ inputs.semantic_min_steps }}"
+
+      - name: Compare baseline vs head_off
+        continue-on-error: true
+        run: |
+          python head/tests/compare_benchmarks.py --baseline baseline.json --current head_off.json --max-rel-diff 0.01 --max-abs-diff 0.0
+
+      - name: Compare head_off vs head_on
+        continue-on-error: true
+        run: |
+          python head/tests/compare_benchmarks.py --baseline head_off.json --current head_on.json --max-rel-diff 0.01 --max-abs-diff 0.0
+
+      - name: Compare baseline vs head_on
+        continue-on-error: true
+        run: |
+          python head/tests/compare_benchmarks.py --baseline baseline.json --current head_on.json --max-rel-diff 0.01 --max-abs-diff 0.0
+
+      - name: Upload benchmark JSON
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: lpips-benchmark-json
+          path: |
+            baseline.json
+            head_off.json
+            head_on.json
diff --git a/tests/benchmark_lpips.py b/tests/benchmark_lpips.py
index d85c499..7ce3163 100644
--- a/tests/benchmark_lpips.py
+++ b/tests/benchmark_lpips.py
@@ -274,6 +274,7 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
     beta_start: float,
     beta_end: float,
     train_timesteps: int,
+    semantic_stop: dict[str, float | int] | None = None,
 ) -> tuple[float, dict[str, float]]:
     import lpips
 
@@ -326,6 +327,10 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
             flow_t = torch.sqrt(1.0 - abt) / (torch.sqrt(1.0 - abt) + torch.sqrt(abt))
             current_times = (sigma, abt, flow_t)
 
+            model_options = {"bench_timestep": int(t)}
+            if semantic_stop is not None:
+                model_options["lanpaint_semantic_stop"] = semantic_stop
+
             x0 = engine(
                 x=x,
                 latent_image=latent_image,
@@ -333,7 +338,7 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
                 sigma=sigma,
                 latent_mask=latent_mask,
                 current_times=current_times,
-                model_options={"bench_timestep": int(t)},
+                model_options=model_options,
                 seed=seed,
                 n_steps=inner_steps,
             )
@@ -380,6 +385,14 @@ def main() -> None:
     parser.add_argument("--beta-start", type=float, default=0.0001, help="DDPM beta_start (guided_diffusion_e2e)")
     parser.add_argument("--beta-end", type=float, default=0.02, help="DDPM beta_end (guided_diffusion_e2e)")
     parser.add_argument("--min-seconds", type=float, default=0.0, help="Fail if benchmark runtime is below this (guided_diffusion_e2e)")
+    parser.add_argument(
+        "--semantic-threshold",
+        type=float,
+        default=0.0,
+        help="Enable LanPaint semantic early-stop if > 0 (guided_diffusion_e2e)",
+    )
+    parser.add_argument("--semantic-patience", type=int, default=1, help="Semantic early-stop patience (guided_diffusion_e2e)")
+    parser.add_argument("--semantic-min-steps", type=int, default=1, help="Semantic early-stop min_steps (guided_diffusion_e2e)")
     args = parser.parse_args()
 
     if args.subset <= 0:
@@ -388,6 +401,10 @@ def main() -> None:
         raise SystemExit("--n-steps must be > 0")
     if args.shuffle_buffer <= 0:
         raise SystemExit("--shuffle-buffer must be > 0")
+    if args.semantic_patience <= 0:
+        raise SystemExit("--semantic-patience must be > 0")
+    if args.semantic_min_steps <= 0:
+        raise SystemExit("--semantic-min-steps must be > 0")
 
     device_str = args.device
     if device_str == "auto":
@@ -423,6 +440,14 @@ def main() -> None:
         if args.lanpaintbench is None or args.checkpoint is None:
             raise SystemExit("--lanpaintbench and --checkpoint are required for guided_diffusion_e2e mode")
 
+        semantic_stop = None
+        if float(args.semantic_threshold) > 0.0:
+            semantic_stop = {
+                "threshold": float(args.semantic_threshold),
+                "patience": int(args.semantic_patience),
+                "min_steps": int(args.semantic_min_steps),
+            }
+
         score, stats = _run_guided_diffusion_e2e_bench(
             LanPaintEngine,
             images,
@@ -436,6 +461,7 @@ def main() -> None:
             beta_start=args.beta_start,
             beta_end=args.beta_end,
             train_timesteps=args.train_timesteps,
+            semantic_stop=semantic_stop,
         )
         if args.min_seconds > 0.0 and stats.get("seconds_total", 0.0) < args.min_seconds:
             raise SystemExit(f"benchmark too fast: {stats.get('seconds_total', 0.0):.1f}s < {args.min_seconds:.1f}s")
@@ -463,6 +489,9 @@ def main() -> None:
                     "lanpaintbench": str(args.lanpaintbench) if args.lanpaintbench is not None else None,
                     "checkpoint": str(args.checkpoint) if args.checkpoint is not None else None,
                     "device": device_str,
+                    "semantic_threshold": args.semantic_threshold,
+                    "semantic_patience": args.semantic_patience,
+                    "semantic_min_steps": args.semantic_min_steps,
                 },
             },
             f,

From 4074473b90223feb518e28617cc3eff150a9593a Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Wed, 14 Jan 2026 23:39:28 +0800
Subject: [PATCH 05/33] Add per-case quality gates to benchmark

---
 .github/workflows/benchmark.yml   |  58 +++++++++++--
 tests/benchmark_lpips.py          |  98 +++++++++++++++++++---
 tests/compare_benchmarks.py       |  82 ++++++++++++++++++
 tests/render_benchmark_gallery.py | 133 ++++++++++++++++++++++++++++++
 4 files changed, 354 insertions(+), 17 deletions(-)
 create mode 100644 tests/render_benchmark_gallery.py

diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml
index e753772..dfa0098 100644
--- a/.github/workflows/benchmark.yml
+++ b/.github/workflows/benchmark.yml
@@ -10,7 +10,7 @@ on:
       subset:
         description: "Number of ImageNet images to benchmark"
         required: true
-        default: "10"
+        default: "1"
       sample_steps:
         description: "Outer diffusion steps (guided_diffusion_e2e)"
         required: true
@@ -43,6 +43,22 @@ on:
         description: "Fail if a benchmark runtime is below this (0 to disable)"
         required: true
         default: "0"
+      max_case_mean_delta:
+        description: "Fail if mean per-case metric delta exceeds this (head_on - head_off)"
+        required: true
+        default: "0.01"
+      max_case_p95_delta:
+        description: "Fail if P95 per-case metric delta exceeds this (head_on - head_off)"
+        required: true
+        default: "0.02"
+      max_case_max_delta:
+        description: "Fail if max per-case metric delta exceeds this (head_on - head_off)"
+        required: true
+        default: "0.05"
+      gallery_topk:
+        description: "How many worst-case examples to render in the gallery"
+        required: true
+        default: "8"
 
 jobs:
   benchmark:
@@ -90,7 +106,7 @@ jobs:
             ~/.cache/huggingface
             ~/.cache/torch
             ~/.cache/lanpaintbench
-          key: ${{ runner.os }}-lpips-imagenet-${{ hashFiles('head/tests/benchmark_lpips.py', 'head/tests/compare_benchmarks.py') }}
+          key: ${{ runner.os }}-lpips-imagenet-${{ hashFiles('head/tests/benchmark_lpips.py', 'head/tests/compare_benchmarks.py', 'head/tests/render_benchmark_gallery.py') }}
 
       - name: Clone LanPaintBench
         run: |
@@ -122,6 +138,7 @@ jobs:
             --subset "${{ inputs.subset }}" \
             --seed "${{ inputs.seed }}" \
             --shuffle-buffer "${{ inputs.shuffle_buffer }}" \
+            --no-shuffle \
             --mask-type box \
             --n-steps "${{ inputs.n_steps }}" \
             --sample-steps "${{ inputs.sample_steps }}" \
@@ -143,12 +160,14 @@ jobs:
             --subset "${{ inputs.subset }}" \
             --seed "${{ inputs.seed }}" \
             --shuffle-buffer "${{ inputs.shuffle_buffer }}" \
+            --no-shuffle \
             --mask-type box \
             --n-steps "${{ inputs.n_steps }}" \
             --sample-steps "${{ inputs.sample_steps }}" \
             --device cpu \
             --min-seconds "${{ inputs.min_seconds }}" \
-            --semantic-threshold 0.0
+            --semantic-threshold 0.0 \
+            --save-images-dir head_off_images
 
       - name: Run Head Benchmark (early-stop ON)
         run: |
@@ -164,6 +183,7 @@ jobs:
             --subset "${{ inputs.subset }}" \
             --seed "${{ inputs.seed }}" \
             --shuffle-buffer "${{ inputs.shuffle_buffer }}" \
+            --no-shuffle \
             --mask-type box \
             --n-steps "${{ inputs.n_steps }}" \
             --sample-steps "${{ inputs.sample_steps }}" \
@@ -171,7 +191,8 @@ jobs:
             --min-seconds "${{ inputs.min_seconds }}" \
             --semantic-threshold "${{ inputs.semantic_threshold }}" \
             --semantic-patience "${{ inputs.semantic_patience }}" \
-            --semantic-min-steps "${{ inputs.semantic_min_steps }}"
+            --semantic-min-steps "${{ inputs.semantic_min_steps }}" \
+            --save-images-dir head_on_images
 
       - name: Compare baseline vs head_off
         continue-on-error: true
@@ -179,21 +200,44 @@ jobs:
           python head/tests/compare_benchmarks.py --baseline baseline.json --current head_off.json --max-rel-diff 0.01 --max-abs-diff 0.0
 
       - name: Compare head_off vs head_on
-        continue-on-error: true
         run: |
-          python head/tests/compare_benchmarks.py --baseline head_off.json --current head_on.json --max-rel-diff 0.01 --max-abs-diff 0.0
+          python head/tests/compare_benchmarks.py \
+            --baseline head_off.json \
+            --current head_on.json \
+            --max-rel-diff 0.01 \
+            --max-abs-diff 0.0 \
+            --case-metric lpips_inpaint_ctx \
+            --max-case-mean-delta "${{ inputs.max_case_mean_delta }}" \
+            --max-case-p95-delta "${{ inputs.max_case_p95_delta }}" \
+            --max-case-max-delta "${{ inputs.max_case_max_delta }}"
 
       - name: Compare baseline vs head_on
         continue-on-error: true
         run: |
           python head/tests/compare_benchmarks.py --baseline baseline.json --current head_on.json --max-rel-diff 0.01 --max-abs-diff 0.0
 
+      - name: Render worst-case gallery (head_off vs head_on)
+        if: ${{ always() }}
+        run: |
+          python head/tests/render_benchmark_gallery.py \
+            --baseline-json head_off.json \
+            --current-json head_on.json \
+            --baseline-images head_off_images \
+            --current-images head_on_images \
+            --out-dir gallery \
+            --metric lpips_inpaint_ctx \
+            --top-k "${{ inputs.gallery_topk }}"
+
       - name: Upload benchmark JSON
         if: ${{ always() }}
         uses: actions/upload-artifact@v4
         with:
-          name: lpips-benchmark-json
+          name: lpips-benchmark-artifacts
           path: |
             baseline.json
             head_off.json
             head_on.json
+            head_off_images
+            head_on_images
+            gallery
+          if-no-files-found: warn
diff --git a/tests/benchmark_lpips.py b/tests/benchmark_lpips.py
index 7ce3163..70ee69d 100644
--- a/tests/benchmark_lpips.py
+++ b/tests/benchmark_lpips.py
@@ -72,14 +72,23 @@ def __call__(self, x, sigma, model_options=None, seed=None):  # type: ignore[no-
         return x0, x0
 
 
-def _load_images(*, dataset: str, split: str, subset: int, seed: int, shuffle_buffer: int) -> list[torch.Tensor]:
+def _load_images(
+    *,
+    dataset: str,
+    split: str,
+    subset: int,
+    seed: int,
+    shuffle_buffer: int,
+    shuffle: bool,
+) -> list[torch.Tensor]:
     from datasets import load_dataset
 
     ds = load_dataset(dataset, split=split, streaming=True)
-    try:
-        ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer)
-    except Exception:
-        pass
+    if shuffle:
+        try:
+            ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer)
+        except Exception:
+            pass
 
     images: list[torch.Tensor] = []
     for idx, item in enumerate(ds):
@@ -150,6 +159,31 @@ def _run_dummy_bench(  # type: ignore[no-untyped-def]
     return float(np.mean(scores))
 
 
+def _tensor_to_pil_rgb(image: torch.Tensor):  # type: ignore[no-untyped-def]
+    from PIL import Image
+
+    if image.ndim == 4:
+        image = image[0]
+    if image.ndim != 3 or image.shape[0] != 3:
+        raise ValueError(f"expected (3,H,W) or (1,3,H,W), got {tuple(image.shape)}")
+    image = image.detach().cpu().clamp(-1.0, 1.0)
+    image = (image + 1.0) * 0.5 * 255.0
+    arr = image.permute(1, 2, 0).to(torch.uint8).numpy()
+    return Image.fromarray(arr, mode="RGB")
+
+
+def _tensor_to_pil_mask(mask: torch.Tensor):  # type: ignore[no-untyped-def]
+    from PIL import Image
+
+    if mask.ndim == 4:
+        mask = mask[0]
+    if mask.ndim != 3:
+        raise ValueError(f"expected (C,H,W) or (1,C,H,W), got {tuple(mask.shape)}")
+    mask_2d = mask[0].detach().cpu().clamp(0.0, 1.0)
+    arr = (mask_2d * 255.0).to(torch.uint8).numpy()
+    return Image.fromarray(arr, mode="L")
+
+
 def _ddpm_alpha_bar(*, timesteps: int, beta_start: float, beta_end: float) -> torch.Tensor:
     betas = torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float64)
     alphas = 1.0 - betas
@@ -275,7 +309,8 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
     beta_end: float,
     train_timesteps: int,
     semantic_stop: dict[str, float | int] | None = None,
-) -> tuple[float, dict[str, float]]:
+    save_images_dir: Path | None = None,
+) -> tuple[float, dict[str, float], list[dict[str, float | int | str]]]:
     import lpips
 
     unet = _load_guided_diffusion_unet(lanpaintbench=lanpaintbench, checkpoint=checkpoint, device=device)
@@ -294,6 +329,7 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
 
     scores: list[float] = []
     start = time.perf_counter()
+    cases: list[dict[str, float | int | str]] = []
 
     engine = LanPaintClass(
         Model=model,
@@ -304,7 +340,11 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
         StepSize=0.1,
     )
 
+    if save_images_dir is not None:
+        save_images_dir.mkdir(parents=True, exist_ok=True)
+
     for idx, img in enumerate(images):
+        case_start = time.perf_counter()
         img = img.unsqueeze(0).to(device)
 
         height, width = img.shape[-2:]
@@ -353,12 +393,43 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
             x = x + d * (sigma_next_ - sigma_)
 
         x = x.clamp(-1.0, 1.0)
-        dist = lpips_fn(x * inpaint_mask, img * inpaint_mask)
-        scores.append(float(dist.item()))
+        reference = img
+
+        # Preserve context by replacing known region with reference in both tensors.
+        x_ctx = x * inpaint_mask + reference * latent_mask
+        ref_ctx = reference
+
+        dist_inpaint_ctx = lpips_fn(x_ctx, ref_ctx)
+        dist_inpaint_zero = lpips_fn(x * inpaint_mask, reference * inpaint_mask)
+        dist_full = lpips_fn(x, reference)
+
+        mse_inpaint = torch.mean(((x - reference) * inpaint_mask) ** 2).item()
+        case_seconds = time.perf_counter() - case_start
+
+        scores.append(float(dist_inpaint_zero.item()))
+        cases.append(
+            {
+                "case_id": idx,
+                "lpips_inpaint_ctx": float(dist_inpaint_ctx.item()),
+                "lpips_inpaint_zero": float(dist_inpaint_zero.item()),
+                "lpips_full": float(dist_full.item()),
+                "mse_inpaint": float(mse_inpaint),
+                "seconds": float(case_seconds),
+            }
+        )
+
+        if save_images_dir is not None:
+            case_dir = save_images_dir / f"{idx:04d}"
+            case_dir.mkdir(parents=True, exist_ok=True)
+            _tensor_to_pil_rgb(reference).save(case_dir / "reference.png")
+            _tensor_to_pil_rgb(latent_image).save(case_dir / "masked_input.png")
+            _tensor_to_pil_rgb(x).save(case_dir / "output.png")
+            _tensor_to_pil_mask(latent_mask).save(case_dir / "known_mask.png")
+            _tensor_to_pil_mask(inpaint_mask).save(case_dir / "inpaint_mask.png")
 
     elapsed = time.perf_counter() - start
     stats = {"seconds_total": elapsed, "seconds_per_image": elapsed / max(1, len(images))}
-    return float(np.mean(scores)), stats
+    return float(np.mean(scores)), stats, cases
 
 
 def main() -> None:
@@ -372,6 +443,7 @@ def main() -> None:
     parser.add_argument("--subset", type=int, default=8)
     parser.add_argument("--seed", type=int, default=42)
     parser.add_argument("--shuffle-buffer", type=int, default=1024)
+    parser.add_argument("--no-shuffle", action="store_true", help="Disable streaming shuffle for deterministic case selection")
 
     parser.add_argument("--mask-type", type=str, default="box", choices=["half", "box", "outpaint", "checkerboard"])
     parser.add_argument("--n-steps", type=int, default=5, help="LanPaint inner iterations per call (dummy mode)")
@@ -393,6 +465,7 @@ def main() -> None:
     )
     parser.add_argument("--semantic-patience", type=int, default=1, help="Semantic early-stop patience (guided_diffusion_e2e)")
     parser.add_argument("--semantic-min-steps", type=int, default=1, help="Semantic early-stop min_steps (guided_diffusion_e2e)")
+    parser.add_argument("--save-images-dir", type=Path, default=None, help="If set, save per-case images for manual review")
     args = parser.parse_args()
 
     if args.subset <= 0:
@@ -423,9 +496,11 @@ def main() -> None:
         subset=args.subset,
         seed=args.seed,
         shuffle_buffer=args.shuffle_buffer,
+        shuffle=not args.no_shuffle,
     )
 
     stats: dict[str, float] = {}
+    cases = None
     if args.mode == "dummy":
         score = _run_dummy_bench(
             LanPaintEngine,
@@ -448,7 +523,7 @@ def main() -> None:
                 "min_steps": int(args.semantic_min_steps),
             }
 
-        score, stats = _run_guided_diffusion_e2e_bench(
+        score, stats, cases = _run_guided_diffusion_e2e_bench(
             LanPaintEngine,
             images,
             device,
@@ -462,6 +537,7 @@ def main() -> None:
             beta_end=args.beta_end,
             train_timesteps=args.train_timesteps,
             semantic_stop=semantic_stop,
+            save_images_dir=args.save_images_dir,
         )
         if args.min_seconds > 0.0 and stats.get("seconds_total", 0.0) < args.min_seconds:
             raise SystemExit(f"benchmark too fast: {stats.get('seconds_total', 0.0):.1f}s < {args.min_seconds:.1f}s")
@@ -479,6 +555,7 @@ def main() -> None:
                     "subset": args.subset,
                     "seed": args.seed,
                     "shuffle_buffer": args.shuffle_buffer,
+                    "no_shuffle": bool(args.no_shuffle),
                     "mask_type": args.mask_type,
                     "n_steps": args.n_steps,
                     "sigma": args.sigma,
@@ -493,6 +570,7 @@ def main() -> None:
                     "semantic_patience": args.semantic_patience,
                     "semantic_min_steps": args.semantic_min_steps,
                 },
+                "cases": cases,
             },
             f,
             indent=2,
diff --git a/tests/compare_benchmarks.py b/tests/compare_benchmarks.py
index f90d595..8861233 100644
--- a/tests/compare_benchmarks.py
+++ b/tests/compare_benchmarks.py
@@ -2,6 +2,7 @@
 
 import argparse
 import json
+import math
 import os
 from pathlib import Path
 
@@ -14,12 +15,34 @@ def _append_step_summary(markdown: str) -> None:
         f.write(markdown)
 
 
+def _quantile(values: list[float], q: float) -> float:
+    if not values:
+        return float("nan")
+    if q <= 0.0:
+        return min(values)
+    if q >= 1.0:
+        return max(values)
+    ordered = sorted(values)
+    # Nearest-rank quantile.
+    idx = max(0, min(len(ordered) - 1, math.ceil(q * len(ordered)) - 1))
+    return float(ordered[idx])
+
+
 def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--baseline", type=Path, required=True)
     parser.add_argument("--current", type=Path, required=True)
     parser.add_argument("--max-rel-diff", type=float, default=0.01, help="Fail if (current-baseline)/baseline exceeds this")
     parser.add_argument("--max-abs-diff", type=float, default=0.0, help="Fail if current-baseline exceeds this")
+    parser.add_argument(
+        "--case-metric",
+        type=str,
+        default="lpips_inpaint_ctx",
+        help="Per-case metric key to compare when 'cases' are present",
+    )
+    parser.add_argument("--max-case-mean-delta", type=float, default=None)
+    parser.add_argument("--max-case-p95-delta", type=float, default=None)
+    parser.add_argument("--max-case-max-delta", type=float, default=None)
     args = parser.parse_args()
 
     baseline = json.loads(args.baseline.read_text(encoding="utf-8"))
@@ -61,6 +84,65 @@ def main() -> None:
     )
     _append_step_summary(summary)
 
+    b_cases = baseline.get("cases") if isinstance(baseline, dict) else None
+    c_cases = current.get("cases") if isinstance(current, dict) else None
+    if isinstance(b_cases, list) and isinstance(c_cases, list) and b_cases and c_cases:
+        b_by_id = {int(row["case_id"]): row for row in b_cases if isinstance(row, dict) and "case_id" in row}
+        c_by_id = {int(row["case_id"]): row for row in c_cases if isinstance(row, dict) and "case_id" in row}
+        shared_ids = sorted(set(b_by_id.keys()) & set(c_by_id.keys()))
+        if shared_ids:
+            deltas: list[float] = []
+            b_vals: list[float] = []
+            c_vals: list[float] = []
+            for case_id in shared_ids:
+                b_row = b_by_id[case_id]
+                c_row = c_by_id[case_id]
+                try:
+                    b_val = float(b_row[args.case_metric])
+                    c_val = float(c_row[args.case_metric])
+                except Exception:
+                    continue
+                b_vals.append(b_val)
+                c_vals.append(c_val)
+                deltas.append(c_val - b_val)
+
+            if deltas:
+                mean_delta = float(sum(deltas) / len(deltas))
+                p95_delta = _quantile(deltas, 0.95)
+                max_delta = max(deltas)
+                b_mean = float(sum(b_vals) / len(b_vals))
+                c_mean = float(sum(c_vals) / len(c_vals))
+
+                print("")
+                print(f"Case metric: {args.case_metric}")
+                print(f"Baseline mean: {b_mean:.6f}")
+                print(f"Current mean:  {c_mean:.6f}")
+                print(f"Mean delta:    {mean_delta:+.6f}")
+                print(f"P95 delta:     {p95_delta:+.6f}")
+                print(f"Max delta:     {max_delta:+.6f}")
+
+                case_summary = "\n".join(
+                    [
+                        "### Per-case deltas",
+                        "",
+                        "| Metric | Mean Δ | P95 Δ | Max Δ | Cases |",
+                        "|---|---:|---:|---:|---:|",
+                        f"| `{args.case_metric}` | {mean_delta:+.6f} | {p95_delta:+.6f} | {max_delta:+.6f} | {len(deltas)} |",
+                        "",
+                    ]
+                )
+                _append_step_summary(case_summary)
+
+                if args.max_case_mean_delta is not None and mean_delta > float(args.max_case_mean_delta):
+                    print("FAIL: per-case mean delta exceeds threshold")
+                    raise SystemExit(1)
+                if args.max_case_p95_delta is not None and p95_delta > float(args.max_case_p95_delta):
+                    print("FAIL: per-case p95 delta exceeds threshold")
+                    raise SystemExit(1)
+                if args.max_case_max_delta is not None and max_delta > float(args.max_case_max_delta):
+                    print("FAIL: per-case max delta exceeds threshold")
+                    raise SystemExit(1)
+
     if diff > args.max_abs_diff:
         if b_score == 0.0 or rel_diff > args.max_rel_diff:
             print("FAIL: LPIPS regression exceeds configured threshold(s)")
diff --git a/tests/render_benchmark_gallery.py b/tests/render_benchmark_gallery.py
new file mode 100644
index 0000000..8671a60
--- /dev/null
+++ b/tests/render_benchmark_gallery.py
@@ -0,0 +1,133 @@
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+
+
+def _load_cases(path: Path) -> dict[int, dict]:
+    data = json.loads(path.read_text(encoding="utf-8"))
+    cases = data.get("cases")
+    if not isinstance(cases, list):
+        return {}
+    out: dict[int, dict] = {}
+    for row in cases:
+        if not isinstance(row, dict) or "case_id" not in row:
+            continue
+        out[int(row["case_id"])] = row
+    return out
+
+
+def _open_rgb(path: Path):  # type: ignore[no-untyped-def]
+    from PIL import Image
+
+    return Image.open(path).convert("RGB")
+
+
+def _open_mask(path: Path):  # type: ignore[no-untyped-def]
+    from PIL import Image
+
+    return Image.open(path).convert("L")
+
+
+def _diff_heatmap(off_rgb, on_rgb, inpaint_mask_l):  # type: ignore[no-untyped-def]
+    from PIL import Image
+
+    off = np.asarray(off_rgb, dtype=np.int16)
+    on = np.asarray(on_rgb, dtype=np.int16)
+    mask = (np.asarray(inpaint_mask_l, dtype=np.uint8) > 127).astype(np.uint8)
+
+    diff = np.abs(off - on).astype(np.uint8)
+    diff = diff.max(axis=2)  # (H,W)
+    diff = diff * mask
+
+    # Scale for visibility; clamp to uint8.
+    diff = np.clip(diff.astype(np.int16) * 4, 0, 255).astype(np.uint8)
+    return Image.fromarray(diff, mode="L").convert("RGB")
+
+
+def _hstack(images):  # type: ignore[no-untyped-def]
+    from PIL import Image
+
+    if not images:
+        raise ValueError("no images to stack")
+    w, h = images[0].size
+    out = Image.new("RGB", (w * len(images), h))
+    for i, im in enumerate(images):
+        out.paste(im, (i * w, 0))
+    return out
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Render a small visual gallery comparing two benchmark runs")
+    parser.add_argument("--baseline-json", type=Path, required=True)
+    parser.add_argument("--current-json", type=Path, required=True)
+    parser.add_argument("--baseline-images", type=Path, required=True)
+    parser.add_argument("--current-images", type=Path, required=True)
+    parser.add_argument("--out-dir", type=Path, required=True)
+    parser.add_argument("--metric", type=str, default="lpips_inpaint_ctx")
+    parser.add_argument("--top-k", type=int, default=8)
+    args = parser.parse_args()
+
+    b_cases = _load_cases(args.baseline_json)
+    c_cases = _load_cases(args.current_json)
+    shared = sorted(set(b_cases.keys()) & set(c_cases.keys()))
+    if not shared:
+        raise SystemExit("no shared cases found in JSON files")
+
+    deltas: list[tuple[float, int]] = []
+    for case_id in shared:
+        try:
+            b_val = float(b_cases[case_id][args.metric])
+            c_val = float(c_cases[case_id][args.metric])
+        except Exception:
+            continue
+        deltas.append((c_val - b_val, case_id))
+
+    if not deltas:
+        raise SystemExit(f"no comparable case metric '{args.metric}' found")
+
+    deltas.sort(reverse=True, key=lambda x: x[0])
+    selected = [case_id for _, case_id in deltas[: max(1, args.top_k)]]
+
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+    rendered: list[str] = []
+
+    for case_id in selected:
+        case_name = f"{case_id:04d}"
+        b_dir = args.baseline_images / case_name
+        c_dir = args.current_images / case_name
+
+        ref = _open_rgb(b_dir / "reference.png")
+        masked = _open_rgb(b_dir / "masked_input.png")
+        off = _open_rgb(b_dir / "output.png")
+        on = _open_rgb(c_dir / "output.png")
+        inpaint_mask = _open_mask(b_dir / "inpaint_mask.png")
+
+        diff = _diff_heatmap(off, on, inpaint_mask)
+        strip = _hstack([ref, masked, off, on, diff])
+
+        out_name = f"{case_name}_compare.png"
+        strip.save(args.out_dir / out_name)
+        rendered.append(out_name)
+
+    index = "\n".join(
+        [
+            "<!doctype html>",
+            "<html><head><meta charset='utf-8'><title>LanPaint benchmark gallery</title></head>",
+            "<body>",
+            "<h2>LanPaint benchmark gallery</h2>",
+            "<p>Columns: reference | masked_input | head_off | head_on | diff(head_off vs head_on)</p>",
+            *[
+                f"<div><img src='{name}' style='image-rendering:pixelated;max-width:100%'></div><hr>"
+                for name in rendered
+            ],
+            "</body></html>",
+            "",
+        ]
+    )
+    (args.out_dir / "index.html").write_text(index, encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()

From 207860edc3f1bd75d91021cfeed3c79b4aa4cdce Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Fri, 16 Jan 2026 16:16:03 +0800
Subject: [PATCH 06/33] Avoid semantic early-stop at low sigma

---
 src/LanPaint/lanpaint.py             |  9 ++++
 tests/test_lanpaint_semantic_stop.py | 75 ++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+)

diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index 7fe94ed..12994c7 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -74,6 +74,15 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
         patience = max(1, patience)
         patience_counter = 0
 
+        if enabled_early_stop:
+            try:
+                abt_val = float(torch.mean(abt).item())
+            except Exception:
+                abt_val = 0.0
+
+            if abt_val > 0.8:
+                min_steps = max(min_steps, int(n_steps))
+
         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 )
 
diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
index 5ed90d2..ce3f45e 100644
--- a/tests/test_lanpaint_semantic_stop.py
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -75,3 +75,78 @@ def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
 
     assert calls["langevin"] == 3
     assert calls["distance"] == 3
+
+
+def test_default_semantic_stop_triggers_at_min_steps_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}
+
+    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
+        return x_t, args
+
+    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "min_steps": 3,
+            "patience": 1,
+        }
+    }
+
+    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
+
+
+def test_semantic_stop_is_disabled_for_very_late_steps() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    calls = {"langevin": 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
+        return x_t, args
+
+    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "min_steps": 2,
+            "patience": 1,
+        }
+    }
+
+    x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
+    current_times = (sigma, torch.tensor([0.99]), torch.tensor([0.0]))
+    engine(
+        x,
+        latent_image,
+        noise,
+        sigma,
+        latent_mask,
+        current_times,
+        model_options=model_options,
+        seed=0,
+        n_steps=10,
+    )
+
+    assert calls["langevin"] == 10

From 896729ccd489d0c86fcb7362af2c0778fdd12d56 Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Fri, 16 Jan 2026 19:01:22 +0800
Subject: [PATCH 07/33] Fix semantic early-stop distance

---
 src/LanPaint/lanpaint.py             |  15 ++--
 tests/test_lanpaint_semantic_stop.py | 103 +++++++++++++++++++++++++++
 2 files changed, 113 insertions(+), 5 deletions(-)

diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index 12994c7..ab4e818 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -80,9 +80,6 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
             except Exception:
                 abt_val = 0.0
 
-            if abt_val > 0.8:
-                min_steps = max(min_steps, int(n_steps))
-
         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 )
 
@@ -109,8 +106,16 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
                         dist = distance_fn(x_t, x_t_prev)
 
                 if dist is None:
-                    diff = (x_t - x_t_prev) * (1 - latent_mask)
-                    dist = torch.mean(diff**2).item()
+                    # Measure masked per-element MSE (normalized by mask area), in fp32 to avoid fp16 quantization.
+                    inpaint = (1 - latent_mask).to(dtype=torch.float32)
+                    diff = (x_t.to(dtype=torch.float32) - x_t_prev.to(dtype=torch.float32)) * inpaint
+                    denom = torch.sum(inpaint) + 1e-12
+                    dist = (torch.sum(diff**2) / denom).item()
+
+                    # Normalize by (1 - abt) because Langevin noise variance scales linearly with (1 - abt).
+                    # This prevents early stopping just because the noise amplitude shrinks at the tail.
+                    scale = max(1.0 - abt_val, 1e-6)
+                    dist = dist / scale
 
                 if float(dist) <= threshold:
                     patience_counter += 1
diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
index ce3f45e..4951631 100644
--- a/tests/test_lanpaint_semantic_stop.py
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -136,6 +136,42 @@ def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y
     }
 
     x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
+
+    # Use a constant per-step update so dist is deterministic, then assert that the
+    # effective threshold becomes stricter as abt -> 1 (so we don't stop too early
+    # in the low-noise regime).
+    delta = 0.01
+
+    def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None):  # type: ignore[no-untyped-def]
+        calls["langevin"] += 1
+        return x_t + delta, args
+
+    engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": delta * delta,
+            "min_steps": 2,
+            "patience": 1,
+        }
+    }
+
+    calls["langevin"] = 0
+    current_times = (sigma, torch.tensor([0.0]), torch.tensor([0.0]))
+    engine(
+        x,
+        latent_image,
+        noise,
+        sigma,
+        latent_mask,
+        current_times,
+        model_options=model_options,
+        seed=0,
+        n_steps=10,
+    )
+    assert calls["langevin"] == 2
+
+    calls["langevin"] = 0
     current_times = (sigma, torch.tensor([0.99]), torch.tensor([0.0]))
     engine(
         x,
@@ -148,5 +184,72 @@ def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y
         seed=0,
         n_steps=10,
     )
+    assert calls["langevin"] == 10
+
+
+def test_default_semantic_stop_is_mask_normalized() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    calls = {"langevin": 0}
+    delta = 0.01
 
+    def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None):  # type: ignore[no-untyped-def]
+        calls["langevin"] += 1
+        return x_t + delta, args
+
+    engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
+
+    # Choose a threshold that would previously depend on how much of the tensor is masked.
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": (delta * delta) * 0.2,
+            "min_steps": 2,
+            "patience": 1,
+        }
+    }
+
+    x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
+    current_times = (sigma, torch.tensor([0.0]), torch.tensor([0.0]))
+
+    # Small hole (unknown region) -> mask mostly 1, with a small 0 patch.
+    small_hole_mask = torch.ones_like(latent_mask)
+    small_hole_mask[:, :, 3:5, 3:5] = 0.0
+
+    calls["langevin"] = 0
+    engine(
+        x,
+        latent_image,
+        noise,
+        sigma,
+        small_hole_mask,
+        current_times,
+        model_options=model_options,
+        seed=0,
+        n_steps=10,
+    )
+    assert calls["langevin"] == 10
+
+    # Larger hole (unknown region).
+    large_hole_mask = torch.ones_like(latent_mask)
+    large_hole_mask[:, :, 1:7, 1:7] = 0.0
+
+    calls["langevin"] = 0
+    engine(
+        x,
+        latent_image,
+        noise,
+        sigma,
+        large_hole_mask,
+        current_times,
+        model_options=model_options,
+        seed=0,
+        n_steps=10,
+    )
     assert calls["langevin"] == 10

From 2086f006037c86c65d34daa561d8b8ed318dcc7c Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 19 Jan 2026 11:26:15 +0800
Subject: [PATCH 08/33] Skip score model after semantic convergence

# Conflicts:
#	src/LanPaint/lanpaint.py
---
 src/LanPaint/lanpaint.py             | 108 +++++++++++++++++++--------
 tests/test_lanpaint_semantic_stop.py |  68 +++++++++++++----
 2 files changed, 130 insertions(+), 46 deletions(-)

diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index ab4e818..bd361e0 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -1,6 +1,7 @@
 import torch
 from .utils import StochasticHarmonicOscillator
 from functools import partial
+import torch.nn.functional as F
 
 class LanPaint():
     def __init__(self, Model, NSteps, Friction, Lambda, Beta, StepSize, IS_FLUX = False, IS_FLOW = False, EarlyStopThreshold = 0.0, EarlyStopPatience = 1, EarlyStopHook = None):
@@ -73,21 +74,39 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
         min_steps = max(1, min_steps)
         patience = max(1, patience)
         patience_counter = 0
+        patience_eff = patience
+        threshold_eff = threshold
+        inpaint_weight = None
+        ring_weight = None
 
         if enabled_early_stop:
             try:
                 abt_val = float(torch.mean(abt).item())
             except Exception:
                 abt_val = 0.0
+            patience_eff = patience + 1 + (1 if abt_val > 0.9 else 0)
+            threshold_scale = max(0.1, (1.0 - abt_val) ** 0.5)
+            threshold_eff = threshold * threshold_scale
 
+            inpaint_weight = (1 - latent_mask).to(dtype=torch.float32)
+            if latent_mask.dim() == 4:
+                mask_f = latent_mask.to(dtype=torch.float32)
+                dilated = F.max_pool2d(mask_f, kernel_size=11, stride=1, padding=5)
+                ring_weight = (dilated - mask_f).clamp(min=0.0, max=1.0) * inpaint_weight
+
+        frozen = False
         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 )
+            score_func = None
+            if not frozen:
+                score_func = partial( self.score_model, y = self.latent_image, mask = latent_mask, abt = self.add_none_dims(abt), sigma = self.add_none_dims(VE_Sigma), tflow = self.add_none_dims(Flow_t), model_options = model_options, seed = seed )
 
-            x_t_prev = x_t.detach().clone() if enabled_early_stop else None
+            prev_args = args
+            x_t_prev = x_t.detach() if enabled_early_stop and (not frozen) and callable(distance_fn) else None
+            x_t_before = x_t if enabled_early_stop and (not frozen) 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 enabled_early_stop and x_t_prev is not None:
+            if enabled_early_stop and (not frozen) and x_t_before is not None:
                 ctx = {
                     "step": i,
                     "steps_done": i + 1,
@@ -106,24 +125,35 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
                         dist = distance_fn(x_t, x_t_prev)
 
                 if dist is None:
-                    # Measure masked per-element MSE (normalized by mask area), in fp32 to avoid fp16 quantization.
-                    inpaint = (1 - latent_mask).to(dtype=torch.float32)
-                    diff = (x_t.to(dtype=torch.float32) - x_t_prev.to(dtype=torch.float32)) * inpaint
-                    denom = torch.sum(inpaint) + 1e-12
-                    dist = (torch.sum(diff**2) / denom).item()
-
-                    # Normalize by (1 - abt) because Langevin noise variance scales linearly with (1 - abt).
-                    # This prevents early stopping just because the noise amplitude shrinks at the tail.
-                    scale = max(1.0 - abt_val, 1e-6)
-                    dist = dist / scale
-
-                if float(dist) <= threshold:
+                    x0_prev = None
+                    x0_cur = None
+                    if isinstance(prev_args, tuple) and len(prev_args) >= 3:
+                        x0_prev = prev_args[2]
+                    if isinstance(args, tuple) and len(args) >= 3:
+                        x0_cur = args[2]
+
+                    if x0_prev is not None and x0_cur is not None:
+                        inpaint = inpaint_weight if inpaint_weight is not None else (1 - latent_mask).to(dtype=torch.float32)
+                        diff_sq = (x0_cur.to(dtype=torch.float32) - x0_prev.to(dtype=torch.float32)) ** 2
+                        denom = torch.sum(inpaint) + 1e-12
+                        dist = (torch.sum(diff_sq * inpaint) / denom).item()
+                        if ring_weight is not None:
+                            ring_denom = torch.sum(ring_weight) + 1e-12
+                            dist_ring = (torch.sum(diff_sq * ring_weight) / ring_denom).item()
+                            dist = max(float(dist), float(dist_ring))
+                    else:
+                        inpaint = inpaint_weight if inpaint_weight is not None else (1 - latent_mask).to(dtype=torch.float32)
+                        diff_sq = (x_t.to(dtype=torch.float32) - x_t_before.to(dtype=torch.float32)) ** 2
+                        denom = torch.sum(inpaint) + 1e-12
+                        dist = (torch.sum(diff_sq * inpaint) / denom).item()
+
+                if float(dist) <= (threshold if callable(distance_fn) else threshold_eff):
                     patience_counter += 1
                 else:
                     patience_counter = 0
 
-                if (i + 1) >= min_steps and patience_counter >= patience:
-                    break
+                if (i + 1) >= min_steps and patience_counter >= patience_eff:
+                    frozen = True
 
         if IS_FLUX or IS_FLOW:
             x = x_t / ( self.add_none_dims(abt)**0.5 + (1-self.add_none_dims(abt))**0.5 )
@@ -173,11 +203,6 @@ def langevin_dynamics(self, x_t, score, mask, step_size, current_times, sigma_x=
         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)
-            C = (abt**0.5 * x0  - x_t )/ (1-abt) + A * x_t
-            return C
         def advance_time(x_t, v, dt, Gamma, A, C, D):
             dtype = x_t.dtype
             with torch.autocast(device_type=x_t.device.type, dtype=torch.float32):
@@ -187,6 +212,23 @@ def advance_time(x_t, v, dt, Gamma, A, C, D):
             v = v.to(dtype)
             return x_t, v
 
+        if score is None:
+            if not (isinstance(args, tuple) and len(args) >= 3):
+                return x_t, args
+            v = args[0]
+            C = args[1]
+            x0 = args[2]
+            x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D)
+            C_new = (abt**0.5 * x0  - x_t )/ (1-abt) + A * x_t
+            v = v + Gamma**0.5 * ( C_new - C) *dt
+            x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C_new, D)
+            return x_t, (v, C_new, x0)
+
+        def Coef_C(x_t):
+            x0 = self.x0_evalutation(x_t, score, sigma, args)
+            C = (abt**0.5 * x0  - x_t )/ (1-abt) + A * x_t
+            return C, x0
+
         def advance_time_overdamped(x_t, dt, A, C, D):
             """
             Overdamped (Gamma -> infinity) limit:
@@ -214,32 +256,32 @@ def advance_time_overdamped(x_t, dt, A, C, D):
         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, C, _ = args
                 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)
+            return x_t, (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
                 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)
+            return x_t, (None, C, x0)
 
         try:
-            x_t_next, (v_next, C_next) = run_damped(x_t, args)
+            x_t_next, (v_next, C_next, x0) = run_damped(x_t, args)
 
             if torch.isnan(x_t_next).any() or torch.isnan(v_next).any():
                 raise ValueError("NaN detected")
@@ -249,9 +291,9 @@ def run_overdamped(x_t, args):
             C = C_next
 
         except Exception:
-            x_t, (v, C) = run_overdamped(x_t, args)
+            x_t, (v, C, x0) = run_overdamped(x_t, args)
 
-        return x_t, (v, C)
+        return x_t, (v, C, x0)
 
     def prepare_step_size(self, current_times, step_size, sigma_x, sigma_y):
         # -------------------------------------------------------------------------
diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
index 4951631..5f0a8eb 100644
--- a/tests/test_lanpaint_semantic_stop.py
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -38,10 +38,14 @@ def test_semantic_stop_triggers_deterministically_at_min_steps() -> None:
         StepSize=0.2,
     )
 
-    calls = {"langevin": 0, "distance": 0}
+    calls = {"langevin": 0, "distance": 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
 
     def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
@@ -73,7 +77,9 @@ def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
         n_steps=10,
     )
 
-    assert calls["langevin"] == 3
+    assert calls["langevin"] == 10
+    assert calls["with_score"] == 3
+    assert calls["without_score"] == 7
     assert calls["distance"] == 3
 
 
@@ -87,10 +93,14 @@ def test_default_semantic_stop_triggers_at_min_steps_without_custom_distance_fn(
         StepSize=0.2,
     )
 
-    calls = {"langevin": 0}
+    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]
@@ -106,7 +116,9 @@ def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y
     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["langevin"] == 10
+    assert calls["with_score"] == 3
+    assert calls["without_score"] == 7
 
 
 def test_semantic_stop_is_disabled_for_very_late_steps() -> None:
@@ -119,10 +131,14 @@ def test_semantic_stop_is_disabled_for_very_late_steps() -> None:
         StepSize=0.2,
     )
 
-    calls = {"langevin": 0}
+    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]
@@ -137,20 +153,25 @@ def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y
 
     x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
 
-    # Use a constant per-step update so dist is deterministic, then assert that the
-    # effective threshold becomes stricter as abt -> 1 (so we don't stop too early
-    # in the low-noise regime).
+    # Use a constant per-step x_t update (simulating Langevin noise) while keeping
+    # an x0-like tensor perfectly stable. Early-stop should ignore the x_t noise.
     delta = 0.01
 
     def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None):  # type: ignore[no-untyped-def]
         calls["langevin"] += 1
-        return x_t + delta, args
+        if score is None:
+            calls["without_score"] += 1
+        else:
+            calls["with_score"] += 1
+        prev_x0 = None if args is None else args[2]
+        x0 = x_t.detach().clone() * 0 if prev_x0 is None else prev_x0
+        return x_t + delta, (None, None, x0)
 
     engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
 
     model_options = {
         "lanpaint_semantic_stop": {
-            "threshold": delta * delta,
+            "threshold": 1e-12,
             "min_steps": 2,
             "patience": 1,
         }
@@ -169,9 +190,13 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         seed=0,
         n_steps=10,
     )
-    assert calls["langevin"] == 2
+    assert calls["langevin"] == 10
+    assert calls["with_score"] == 3
+    assert calls["without_score"] == 7
 
     calls["langevin"] = 0
+    calls["with_score"] = 0
+    calls["without_score"] = 0
     current_times = (sigma, torch.tensor([0.99]), torch.tensor([0.0]))
     engine(
         x,
@@ -185,6 +210,8 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         n_steps=10,
     )
     assert calls["langevin"] == 10
+    assert calls["with_score"] == 4
+    assert calls["without_score"] == 6
 
 
 def test_default_semantic_stop_is_mask_normalized() -> None:
@@ -197,12 +224,19 @@ def test_default_semantic_stop_is_mask_normalized() -> None:
         StepSize=0.2,
     )
 
-    calls = {"langevin": 0}
+    calls = {"langevin": 0, "with_score": 0, "without_score": 0}
     delta = 0.01
 
     def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None):  # type: ignore[no-untyped-def]
         calls["langevin"] += 1
-        return x_t + delta, args
+        if score is None:
+            calls["without_score"] += 1
+        else:
+            calls["with_score"] += 1
+        prev_x0 = None if args is None else args[2]
+        x0 = x_t.detach().clone() * 0 if prev_x0 is None else prev_x0
+        x0 = x0 + delta * (1 - mask)
+        return x_t, (None, None, x0)
 
     engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
 
@@ -223,6 +257,8 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
     small_hole_mask[:, :, 3:5, 3:5] = 0.0
 
     calls["langevin"] = 0
+    calls["with_score"] = 0
+    calls["without_score"] = 0
     engine(
         x,
         latent_image,
@@ -235,12 +271,16 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         n_steps=10,
     )
     assert calls["langevin"] == 10
+    assert calls["with_score"] == 10
+    assert calls["without_score"] == 0
 
     # Larger hole (unknown region).
     large_hole_mask = torch.ones_like(latent_mask)
     large_hole_mask[:, :, 1:7, 1:7] = 0.0
 
     calls["langevin"] = 0
+    calls["with_score"] = 0
+    calls["without_score"] = 0
     engine(
         x,
         latent_image,
@@ -253,3 +293,5 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         n_steps=10,
     )
     assert calls["langevin"] == 10
+    assert calls["with_score"] == 10
+    assert calls["without_score"] == 0

From e7080f9574dadb6ed089a5e4189086bbfb07d83b Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Sat, 17 Jan 2026 11:18:31 +0800
Subject: [PATCH 09/33] Test ring changes prevent premature freeze

---
 tests/test_lanpaint_semantic_stop.py | 89 ++++++++++++++++++++++++++++
 1 file changed, 89 insertions(+)

diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
index 5f0a8eb..313876b 100644
--- a/tests/test_lanpaint_semantic_stop.py
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -1,4 +1,5 @@
 import torch
+import torch.nn.functional as F
 
 from src.LanPaint.lanpaint import LanPaint as LanPaintEngine
 
@@ -295,3 +296,91 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
     assert calls["langevin"] == 10
     assert calls["with_score"] == 10
     assert calls["without_score"] == 0
+
+
+def test_semantic_stop_does_not_freeze_while_ring_changes() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    batch = 1
+    channels = 4
+    height = 64
+    width = 64
+
+    x = torch.zeros((batch, channels, height, width))
+    latent_image = torch.zeros_like(x)
+    noise = torch.ones_like(x)
+    sigma = torch.tensor([1.0])
+
+    # latent_mask uses LanPaint convention: 1=known region, 0=to-inpaint.
+    latent_mask = torch.ones_like(x)
+    latent_mask[:, :, 12:52, 12:52] = 0.0
+    inpaint_mask = 1.0 - latent_mask
+
+    # Replicate the "ring" construction in LanPaint: a dilated mask boundary
+    # intersected with the inpaint region.
+    mask_f = latent_mask.to(dtype=torch.float32)
+    dilated = F.max_pool2d(mask_f, kernel_size=11, stride=1, padding=5)
+    ring_mask = (dilated - mask_f).clamp(min=0.0, max=1.0) * inpaint_mask.to(dtype=torch.float32)
+
+    ring_area = float(torch.sum(ring_mask).item())
+    inpaint_area = float(torch.sum(inpaint_mask).item())
+    assert ring_area > 0.0
+    assert ring_area < inpaint_area
+
+    threshold = 1e-6
+    ratio = inpaint_area / ring_area
+    delta_sq = threshold * ratio * 0.9
+    delta = float(delta_sq**0.5)
+
+    calls = {"langevin": 0, "with_score": 0, "without_score": 0}
+
+    def fake_langevin_delta(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
+
+        # Simulate a case where the interior is stable, but the boundary ring keeps changing.
+        step = float(calls["langevin"])
+        x0 = ring_mask.to(dtype=x_t.dtype) * (step * delta)
+        return x_t, (None, None, x0)
+
+    engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
+
+    current_times = (sigma, torch.tensor([0.0]), torch.tensor([0.0]))
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": threshold,
+            "min_steps": 2,
+            "patience": 1,
+        }
+    }
+
+    engine(
+        x,
+        latent_image,
+        noise,
+        sigma,
+        latent_mask,
+        current_times,
+        model_options=model_options,
+        seed=0,
+        n_steps=10,
+    )
+
+    # If ring changes are respected, we should never "freeze" (i.e. score never becomes None).
+    assert calls["langevin"] == 10
+    assert calls["with_score"] == 10
+    assert calls["without_score"] == 0
+
+    # If ring changes were ignored, inpaint-average MSE would fall below the threshold
+    # and the sampler would freeze early. This test ensures we keep running the model
+    # while the boundary is still changing.

From 450a25e5e8ab1f23358ab2641ade67871766e6a4 Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Sun, 18 Jan 2026 17:23:54 +0800
Subject: [PATCH 10/33] Refine semantic inner stop

---
 src/LanPaint/lanpaint.py             | 109 +++++++++++++++++++++------
 tests/benchmark_lpips.py             |  78 ++++++++++++++++++-
 tests/test_lanpaint_semantic_stop.py |  73 +++++++++++++++---
 3 files changed, 222 insertions(+), 38 deletions(-)

diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index bd361e0..9cdcd34 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -78,35 +78,42 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
         threshold_eff = threshold
         inpaint_weight = None
         ring_weight = None
+        trace = None
 
         if enabled_early_stop:
             try:
                 abt_val = float(torch.mean(abt).item())
             except Exception:
                 abt_val = 0.0
-            patience_eff = patience + 1 + (1 if abt_val > 0.9 else 0)
-            threshold_scale = max(0.1, (1.0 - abt_val) ** 0.5)
-            threshold_eff = threshold * threshold_scale
 
-            inpaint_weight = (1 - latent_mask).to(dtype=torch.float32)
-            if latent_mask.dim() == 4:
-                mask_f = latent_mask.to(dtype=torch.float32)
-                dilated = F.max_pool2d(mask_f, kernel_size=11, stride=1, padding=5)
-                ring_weight = (dilated - mask_f).clamp(min=0.0, max=1.0) * inpaint_weight
-
-        frozen = False
+            # Skip semantic early-stop in extremely noisy steps (low abt) and
+            # at the extreme tail where even tiny changes can matter.
+            if abt_val < 0.15 or abt_val > 0.999:
+                enabled_early_stop = False
+            else:
+                patience_eff = patience + 1 + (1 if abt_val > 0.9 else 0)
+                threshold_scale = max(0.1, (1.0 - abt_val) ** 0.5)
+                threshold_eff = threshold * threshold_scale
+
+                inpaint_weight = (1 - latent_mask).to(dtype=torch.float32)
+                if latent_mask.dim() == 4:
+                    mask_f = latent_mask.to(dtype=torch.float32)
+                    dilated = F.max_pool2d(mask_f, kernel_size=11, stride=1, padding=5)
+                    ring_weight = (dilated - mask_f).clamp(min=0.0, max=1.0) * inpaint_weight
+                if isinstance(model_options, dict):
+                    trace = model_options.get("lanpaint_semantic_trace")
+
+        x0_anchor = None
         for i in range(n_steps):
-            score_func = None
-            if not frozen:
-                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 )
+            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 enabled_early_stop and (not frozen) and callable(distance_fn) else None
-            x_t_before = x_t if enabled_early_stop and (not frozen) else None
+            x_t_prev = x_t.detach() if enabled_early_stop and callable(distance_fn) else None
+            x_t_before = x_t if enabled_early_stop 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 enabled_early_stop and (not frozen) and x_t_before is not None:
+            if enabled_early_stop and x_t_before is not None:
                 ctx = {
                     "step": i,
                     "steps_done": i + 1,
@@ -118,15 +125,21 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
                 }
 
                 dist = None
+                custom_dist = False
+                dist_inpaint = None
+                dist_ring = None
+                dist_drift = None
+                x0_prev = None
+                x0_cur = None
+
                 if callable(distance_fn):
                     try:
                         dist = distance_fn(x_t_prev, x_t, ctx)
                     except TypeError:
                         dist = distance_fn(x_t, x_t_prev)
+                    custom_dist = dist is not None
 
                 if dist is None:
-                    x0_prev = None
-                    x0_cur = None
                     if isinstance(prev_args, tuple) and len(prev_args) >= 3:
                         x0_prev = prev_args[2]
                     if isinstance(args, tuple) and len(args) >= 3:
@@ -136,24 +149,72 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
                         inpaint = inpaint_weight if inpaint_weight is not None else (1 - latent_mask).to(dtype=torch.float32)
                         diff_sq = (x0_cur.to(dtype=torch.float32) - x0_prev.to(dtype=torch.float32)) ** 2
                         denom = torch.sum(inpaint) + 1e-12
-                        dist = (torch.sum(diff_sq * inpaint) / denom).item()
+                        dist_inpaint = (torch.sum(diff_sq * inpaint) / denom).item()
+                        dist = float(dist_inpaint)
                         if ring_weight is not None:
                             ring_denom = torch.sum(ring_weight) + 1e-12
                             dist_ring = (torch.sum(diff_sq * ring_weight) / ring_denom).item()
-                            dist = max(float(dist), float(dist_ring))
+                            dist = max(float(dist_inpaint), float(dist_ring))
                     else:
                         inpaint = inpaint_weight if inpaint_weight is not None else (1 - latent_mask).to(dtype=torch.float32)
                         diff_sq = (x_t.to(dtype=torch.float32) - x_t_before.to(dtype=torch.float32)) ** 2
                         denom = torch.sum(inpaint) + 1e-12
                         dist = (torch.sum(diff_sq * inpaint) / denom).item()
+                        dist_inpaint = dist
+
+                threshold_used = threshold if custom_dist else threshold_eff
+
+                # Extra guard: even if per-step changes are tiny, avoid stopping if the
+                # x0 estimate keeps drifting over several steps.
+                if x0_cur is not None and not custom_dist:
+                    if float(dist) <= threshold_used:
+                        if x0_anchor is None:
+                            x0_anchor = x0_cur.detach()
+                        else:
+                            inpaint = inpaint_weight if inpaint_weight is not None else (1 - latent_mask).to(dtype=torch.float32)
+                            diff_sq = (x0_cur.to(dtype=torch.float32) - x0_anchor.to(dtype=torch.float32)) ** 2
+                            denom = torch.sum(inpaint) + 1e-12
+                            drift_inpaint = (torch.sum(diff_sq * inpaint) / denom).item()
+                            dist_drift = float(drift_inpaint)
+                            if ring_weight is not None:
+                                ring_denom = torch.sum(ring_weight) + 1e-12
+                                drift_ring = (torch.sum(diff_sq * ring_weight) / ring_denom).item()
+                                dist_drift = max(float(drift_inpaint), float(drift_ring))
+                            dist = max(float(dist), float(dist_drift))
+                    else:
+                        x0_anchor = None
 
-                if float(dist) <= (threshold if callable(distance_fn) else threshold_eff):
+                if float(dist) <= threshold_used:
                     patience_counter += 1
                 else:
                     patience_counter = 0
-
-                if (i + 1) >= min_steps and patience_counter >= patience_eff:
-                    frozen = True
+                    x0_anchor = None
+
+                should_stop = (i + 1) >= min_steps and patience_counter >= patience_eff
+                if isinstance(trace, list):
+                    trace.append(
+                        {
+                            "case_id": model_options.get("bench_case_id") if isinstance(model_options, dict) else None,
+                            "outer_step": model_options.get("bench_outer_step") if isinstance(model_options, dict) else None,
+                            "bench_timestep": model_options.get("bench_timestep") if isinstance(model_options, dict) else None,
+                            "inner_step": i + 1,
+                            "dist": float(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(threshold_eff),
+                            "patience_counter": int(patience_counter),
+                            "patience_eff": int(patience_eff),
+                            "min_steps": int(min_steps),
+                            "abt": float(abt_val),
+                            "custom_dist": bool(custom_dist),
+                            "stopped": bool(should_stop),
+                        }
+                    )
+
+                if should_stop:
+                    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 )
diff --git a/tests/benchmark_lpips.py b/tests/benchmark_lpips.py
index 70ee69d..8d2a0f7 100644
--- a/tests/benchmark_lpips.py
+++ b/tests/benchmark_lpips.py
@@ -225,9 +225,11 @@ def __init__(self, *, unet, device: torch.device) -> None:  # type: ignore[no-un
         self.model_sampling = _GuidedDiffusionSampling()
         self._unet = unet
         self._device = device
+        self.call_count = 0
 
     @torch.no_grad()
     def __call__(self, x, sigma, model_options=None, seed=None):  # type: ignore[no-untyped-def]
+        self.call_count += 1
         if not isinstance(model_options, dict) or "bench_timestep" not in model_options:
             raise ValueError("model_options['bench_timestep'] is required for guided_diffusion mode")
 
@@ -309,6 +311,9 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
     beta_end: float,
     train_timesteps: int,
     semantic_stop: dict[str, float | int] | None = None,
+    outer_cutoff_steps: int = 0,
+    semantic_trace: list[dict[str, float | int | str | None]] | None = None,
+    report_earlystop: bool = False,
     save_images_dir: Path | None = None,
 ) -> tuple[float, dict[str, float], list[dict[str, float | int | str]]]:
     import lpips
@@ -331,6 +336,14 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
     start = time.perf_counter()
     cases: list[dict[str, float | int | str]] = []
 
+    outer_steps = int(sigmas.shape[0] - 1)
+    n_images = max(1, len(images))
+    model_calls_total = 0
+    inner_steps_executed_total = 0
+    inner_steps_requested_total = 0
+    inner_steps_executed_by_outer_step = [0] * outer_steps
+    inner_steps_requested_by_outer_step = [0] * outer_steps
+
     engine = LanPaintClass(
         Model=model,
         NSteps=inner_steps,
@@ -358,6 +371,9 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
 
         x = latent_image + noise * sigmas[0].reshape((1,) + (1,) * (img.ndim - 1))
 
+        case_model_calls = 0
+        case_inner_steps_executed = 0
+        case_inner_steps_requested = 0
         for step_idx in range(sigmas.shape[0] - 1):
             sigma = sigmas[step_idx : step_idx + 1]
             sigma_next = sigmas[step_idx + 1 : step_idx + 2]
@@ -367,10 +383,22 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
             flow_t = torch.sqrt(1.0 - abt) / (torch.sqrt(1.0 - abt) + torch.sqrt(abt))
             current_times = (sigma, abt, flow_t)
 
-            model_options = {"bench_timestep": int(t)}
+            remaining = outer_steps - step_idx
+            inner_steps_eff = inner_steps
+            if outer_cutoff_steps > 0 and remaining <= outer_cutoff_steps:
+                inner_steps_eff = 0
+
+            model_options: dict[str, object] = {
+                "bench_timestep": int(t),
+                "bench_case_id": int(idx),
+                "bench_outer_step": int(step_idx),
+            }
             if semantic_stop is not None:
                 model_options["lanpaint_semantic_stop"] = semantic_stop
+            if semantic_trace is not None:
+                model_options["lanpaint_semantic_trace"] = semantic_trace
 
+            model.call_count = 0
             x0 = engine(
                 x=x,
                 latent_image=latent_image,
@@ -380,8 +408,20 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
                 current_times=current_times,
                 model_options=model_options,
                 seed=seed,
-                n_steps=inner_steps,
+                n_steps=inner_steps_eff,
             )
+            calls_this_step = int(model.call_count)
+            case_model_calls += calls_this_step
+            model_calls_total += calls_this_step
+
+            executed = max(0, calls_this_step - 1)
+            case_inner_steps_executed += executed
+            inner_steps_executed_total += executed
+            inner_steps_executed_by_outer_step[step_idx] += executed
+
+            case_inner_steps_requested += int(inner_steps_eff)
+            inner_steps_requested_total += int(inner_steps_eff)
+            inner_steps_requested_by_outer_step[step_idx] += int(inner_steps_eff)
 
             if float(sigma_next.item()) == 0.0:
                 x = x0
@@ -414,6 +454,9 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
                 "lpips_inpaint_zero": float(dist_inpaint_zero.item()),
                 "lpips_full": float(dist_full.item()),
                 "mse_inpaint": float(mse_inpaint),
+                "model_calls": int(case_model_calls),
+                "inner_steps_executed": int(case_inner_steps_executed),
+                "inner_steps_requested": int(case_inner_steps_requested),
                 "seconds": float(case_seconds),
             }
         )
@@ -429,6 +472,17 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
 
     elapsed = time.perf_counter() - start
     stats = {"seconds_total": elapsed, "seconds_per_image": elapsed / max(1, len(images))}
+    if report_earlystop:
+        stats["earlystop"] = {
+            "model_calls_total": int(model_calls_total),
+            "model_calls_per_image": float(model_calls_total / n_images),
+            "inner_steps_executed_total": int(inner_steps_executed_total),
+            "inner_steps_executed_per_image": float(inner_steps_executed_total / n_images),
+            "inner_steps_requested_total": int(inner_steps_requested_total),
+            "inner_steps_requested_per_image": float(inner_steps_requested_total / n_images),
+            "inner_steps_executed_by_outer_step_mean": [float(v / n_images) for v in inner_steps_executed_by_outer_step],
+            "inner_steps_requested_by_outer_step_mean": [float(v / n_images) for v in inner_steps_requested_by_outer_step],
+        }
     return float(np.mean(scores)), stats, cases
 
 
@@ -457,6 +511,12 @@ def main() -> None:
     parser.add_argument("--beta-start", type=float, default=0.0001, help="DDPM beta_start (guided_diffusion_e2e)")
     parser.add_argument("--beta-end", type=float, default=0.02, help="DDPM beta_end (guided_diffusion_e2e)")
     parser.add_argument("--min-seconds", type=float, default=0.0, help="Fail if benchmark runtime is below this (guided_diffusion_e2e)")
+    parser.add_argument(
+        "--outer-cutoff-steps",
+        type=int,
+        default=0,
+        help="Disable LanPaint inner iterations for the last N outer steps (guided_diffusion_e2e)",
+    )
     parser.add_argument(
         "--semantic-threshold",
         type=float,
@@ -465,6 +525,8 @@ def main() -> None:
     )
     parser.add_argument("--semantic-patience", type=int, default=1, help="Semantic early-stop patience (guided_diffusion_e2e)")
     parser.add_argument("--semantic-min-steps", type=int, default=1, help="Semantic early-stop min_steps (guided_diffusion_e2e)")
+    parser.add_argument("--report-earlystop", action="store_true", help="Include early-stop stats in output JSON (guided_diffusion_e2e)")
+    parser.add_argument("--trace-inner", type=Path, default=None, help="If set, write per-inner-step semantic trace JSON (guided_diffusion_e2e)")
     parser.add_argument("--save-images-dir", type=Path, default=None, help="If set, save per-case images for manual review")
     args = parser.parse_args()
 
@@ -478,6 +540,8 @@ def main() -> None:
         raise SystemExit("--semantic-patience must be > 0")
     if args.semantic_min_steps <= 0:
         raise SystemExit("--semantic-min-steps must be > 0")
+    if args.outer_cutoff_steps < 0:
+        raise SystemExit("--outer-cutoff-steps must be >= 0")
 
     device_str = args.device
     if device_str == "auto":
@@ -523,6 +587,7 @@ def main() -> None:
                 "min_steps": int(args.semantic_min_steps),
             }
 
+        semantic_trace = [] if args.trace_inner is not None else None
         score, stats, cases = _run_guided_diffusion_e2e_bench(
             LanPaintEngine,
             images,
@@ -537,10 +602,16 @@ def main() -> None:
             beta_end=args.beta_end,
             train_timesteps=args.train_timesteps,
             semantic_stop=semantic_stop,
+            outer_cutoff_steps=int(args.outer_cutoff_steps),
+            semantic_trace=semantic_trace,
+            report_earlystop=bool(args.report_earlystop),
             save_images_dir=args.save_images_dir,
         )
         if args.min_seconds > 0.0 and stats.get("seconds_total", 0.0) < args.min_seconds:
             raise SystemExit(f"benchmark too fast: {stats.get('seconds_total', 0.0):.1f}s < {args.min_seconds:.1f}s")
+        if args.trace_inner is not None and semantic_trace is not None:
+            args.trace_inner.parent.mkdir(parents=True, exist_ok=True)
+            args.trace_inner.write_text(json.dumps(semantic_trace), encoding="utf-8")
 
     args.output.parent.mkdir(parents=True, exist_ok=True)
     with args.output.open("w", encoding="utf-8") as f:
@@ -569,6 +640,9 @@ def main() -> None:
                     "semantic_threshold": args.semantic_threshold,
                     "semantic_patience": args.semantic_patience,
                     "semantic_min_steps": args.semantic_min_steps,
+                    "outer_cutoff_steps": args.outer_cutoff_steps,
+                    "report_earlystop": bool(args.report_earlystop),
+                    "trace_inner": str(args.trace_inner) if args.trace_inner is not None else None,
                 },
                 "cases": cases,
             },
diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
index 313876b..88451c6 100644
--- a/tests/test_lanpaint_semantic_stop.py
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -78,9 +78,9 @@ def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
         n_steps=10,
     )
 
-    assert calls["langevin"] == 10
+    assert calls["langevin"] == 3
     assert calls["with_score"] == 3
-    assert calls["without_score"] == 7
+    assert calls["without_score"] == 0
     assert calls["distance"] == 3
 
 
@@ -117,9 +117,9 @@ def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y
     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"] == 10
+    assert calls["langevin"] == 3
     assert calls["with_score"] == 3
-    assert calls["without_score"] == 7
+    assert calls["without_score"] == 0
 
 
 def test_semantic_stop_is_disabled_for_very_late_steps() -> None:
@@ -179,7 +179,7 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
     }
 
     calls["langevin"] = 0
-    current_times = (sigma, torch.tensor([0.0]), torch.tensor([0.0]))
+    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
     engine(
         x,
         latent_image,
@@ -191,9 +191,9 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         seed=0,
         n_steps=10,
     )
-    assert calls["langevin"] == 10
+    assert calls["langevin"] == 3
     assert calls["with_score"] == 3
-    assert calls["without_score"] == 7
+    assert calls["without_score"] == 0
 
     calls["langevin"] = 0
     calls["with_score"] = 0
@@ -210,9 +210,58 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         seed=0,
         n_steps=10,
     )
-    assert calls["langevin"] == 10
+    assert calls["langevin"] == 4
     assert calls["with_score"] == 4
-    assert calls["without_score"] == 6
+    assert calls["without_score"] == 0
+
+
+def test_semantic_stop_is_disabled_for_very_early_steps() -> 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,
+            "min_steps": 2,
+            "patience": 1,
+        }
+    }
+
+    x, latent_image, noise, sigma, latent_mask, _ = _inputs()
+    current_times = (sigma, torch.tensor([0.0]), torch.tensor([0.0]))
+    engine(
+        x,
+        latent_image,
+        noise,
+        sigma,
+        latent_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
 
 
 def test_default_semantic_stop_is_mask_normalized() -> None:
@@ -298,7 +347,7 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
     assert calls["without_score"] == 0
 
 
-def test_semantic_stop_does_not_freeze_while_ring_changes() -> None:
+def test_semantic_stop_does_not_stop_while_ring_changes() -> None:
     engine = LanPaintEngine(
         _DummyModel(),
         NSteps=10,
@@ -376,11 +425,11 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         n_steps=10,
     )
 
-    # If ring changes are respected, we should never "freeze" (i.e. score never becomes None).
+    # If ring changes are respected, we should never stop early.
     assert calls["langevin"] == 10
     assert calls["with_score"] == 10
     assert calls["without_score"] == 0
 
     # If ring changes were ignored, inpaint-average MSE would fall below the threshold
-    # and the sampler would freeze early. This test ensures we keep running the model
+    # and the sampler would stop early. This test ensures we keep running the model
     # while the boundary is still changing.

From 9c9b5555fc198ef47ce458556d35fd84c3cb5d6c Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 19 Jan 2026 00:41:35 +0800
Subject: [PATCH 11/33] Tighten semantic stop in mid-noise steps

---
 src/LanPaint/lanpaint.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index 9cdcd34..8b2761c 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -91,7 +91,13 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
             if abt_val < 0.15 or abt_val > 0.999:
                 enabled_early_stop = False
             else:
-                patience_eff = patience + 1 + (1 if abt_val > 0.9 else 0)
+                # More noise -> require more consecutive stable steps before stopping.
+                # This keeps early-stop conservative in mid-noise outer steps.
+                patience_eff = patience + 1
+                if abt_val < 0.5:
+                    patience_eff += 1
+                if abt_val > 0.9:
+                    patience_eff += 1
                 threshold_scale = max(0.1, (1.0 - abt_val) ** 0.5)
                 threshold_eff = threshold * threshold_scale
 

From 6fd5f603838103822bc67f618581ef21460c6188 Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 19 Jan 2026 11:29:38 +0800
Subject: [PATCH 12/33] Adjust SHO regression test for x0 return

---
 tests/test_sho_regression.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/tests/test_sho_regression.py b/tests/test_sho_regression.py
index a601532..02102ae 100644
--- a/tests/test_sho_regression.py
+++ b/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,13 +28,16 @@ def score(x):
         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)
+        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)
@@ -47,5 +52,7 @@ def test_langevin_dynamics_fallback_on_exception() -> None:
         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()
\ No newline at end of file
+        x_out, args_out = lp.langevin_dynamics(x_t, score, mask, step_size, current_times, sigma_y=1.0)
+        v_out = args_out[0]
+        assert torch.isfinite(x_out).all()
+        assert v_out is None or torch.isfinite(v_out).all()

From 73ab537e86660eb25341b97674e6202f0c25141d Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 19 Jan 2026 11:37:24 +0800
Subject: [PATCH 13/33] Make CI-only imports lighter

---
 src/LanPaint/lanpaint.py | 3 ++-
 tests/benchmark_lpips.py | 7 +++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index 8b2761c..df81cb3 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -1,7 +1,6 @@
 import torch
 from .utils import StochasticHarmonicOscillator
 from functools import partial
-import torch.nn.functional as F
 
 class LanPaint():
     def __init__(self, Model, NSteps, Friction, Lambda, Beta, StepSize, IS_FLUX = False, IS_FLOW = False, EarlyStopThreshold = 0.0, EarlyStopPatience = 1, EarlyStopHook = None):
@@ -103,6 +102,8 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
 
                 inpaint_weight = (1 - latent_mask).to(dtype=torch.float32)
                 if latent_mask.dim() == 4:
+                    from torch.nn import functional as F
+
                     mask_f = latent_mask.to(dtype=torch.float32)
                     dilated = F.max_pool2d(mask_f, kernel_size=11, stride=1, padding=5)
                     ring_weight = (dilated - mask_f).clamp(min=0.0, max=1.0) * inpaint_weight
diff --git a/tests/benchmark_lpips.py b/tests/benchmark_lpips.py
index 8d2a0f7..75fc297 100644
--- a/tests/benchmark_lpips.py
+++ b/tests/benchmark_lpips.py
@@ -11,7 +11,6 @@
 import time
 from pathlib import Path
 
-import numpy as np
 import torch
 
 
@@ -81,6 +80,7 @@ def _load_images(
     shuffle_buffer: int,
     shuffle: bool,
 ) -> list[torch.Tensor]:
+    import numpy as np
     from datasets import load_dataset
 
     ds = load_dataset(dataset, split=split, streaming=True)
@@ -156,7 +156,7 @@ def _run_dummy_bench(  # type: ignore[no-untyped-def]
         dist = lpips_fn(output, img)
         scores.append(float(dist.item()))
 
-    return float(np.mean(scores))
+    return float(sum(scores) / max(1, len(scores)))
 
 
 def _tensor_to_pil_rgb(image: torch.Tensor):  # type: ignore[no-untyped-def]
@@ -483,7 +483,7 @@ def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
             "inner_steps_executed_by_outer_step_mean": [float(v / n_images) for v in inner_steps_executed_by_outer_step],
             "inner_steps_requested_by_outer_step_mean": [float(v / n_images) for v in inner_steps_requested_by_outer_step],
         }
-    return float(np.mean(scores)), stats, cases
+    return float(sum(scores) / max(1, len(scores))), stats, cases
 
 
 def main() -> None:
@@ -552,7 +552,6 @@ def main() -> None:
     from LanPaint.lanpaint import LanPaint as LanPaintEngine
 
     torch.manual_seed(args.seed)
-    np.random.seed(args.seed)
 
     images = _load_images(
         dataset=args.dataset,

From ac7a5a1f2f35f6f9b348d2fa7ac43d71a4fb5de1 Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 19 Jan 2026 16:42:39 +0800
Subject: [PATCH 14/33] Fix semantic-stop distance_fn dispatch

---
 src/LanPaint/lanpaint.py             | 100 ++++++++++++---
 tests/benchmark_lpips.py             |   5 +-
 tests/compare_benchmarks.py          |   4 +-
 tests/render_benchmark_gallery.py    |   4 +-
 tests/test_lanpaint_semantic_stop.py | 180 +++++++++++++++++++++++++++
 5 files changed, 273 insertions(+), 20 deletions(-)

diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index df81cb3..db4e265 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -1,6 +1,16 @@
 import torch
 from .utils import StochasticHarmonicOscillator
 from functools import partial
+import inspect
+
+# Early-stop constants
+STOP_THRESHOLD_MIN_ABT = 0.15
+STOP_THRESHOLD_MAX_ABT = 0.999
+PATIENCE_BOOST_LOW_ABT = 0.5
+PATIENCE_BOOST_HIGH_ABT = 0.9
+THRESHOLD_SCALE_MIN = 0.1
+RING_KERNEL_SIZE = 11
+RING_PADDING = 5
 
 class LanPaint():
     def __init__(self, Model, NSteps, Friction, Lambda, Beta, StepSize, IS_FLUX = False, IS_FLOW = False, EarlyStopThreshold = 0.0, EarlyStopPatience = 1, EarlyStopHook = None):
@@ -78,6 +88,7 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
         inpaint_weight = None
         ring_weight = None
         trace = None
+        dist_wrapper = None
 
         if enabled_early_stop:
             try:
@@ -87,17 +98,17 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
 
             # Skip semantic early-stop in extremely noisy steps (low abt) and
             # at the extreme tail where even tiny changes can matter.
-            if abt_val < 0.15 or abt_val > 0.999:
+            if abt_val < STOP_THRESHOLD_MIN_ABT or abt_val > STOP_THRESHOLD_MAX_ABT:
                 enabled_early_stop = False
             else:
                 # More noise -> require more consecutive stable steps before stopping.
                 # This keeps early-stop conservative in mid-noise outer steps.
                 patience_eff = patience + 1
-                if abt_val < 0.5:
+                if abt_val < PATIENCE_BOOST_LOW_ABT:
                     patience_eff += 1
-                if abt_val > 0.9:
+                if abt_val > PATIENCE_BOOST_HIGH_ABT:
                     patience_eff += 1
-                threshold_scale = max(0.1, (1.0 - abt_val) ** 0.5)
+                threshold_scale = max(THRESHOLD_SCALE_MIN, (1.0 - abt_val) ** 0.5)
                 threshold_eff = threshold * threshold_scale
 
                 inpaint_weight = (1 - latent_mask).to(dtype=torch.float32)
@@ -105,12 +116,71 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
                     from torch.nn import functional as F
 
                     mask_f = latent_mask.to(dtype=torch.float32)
-                    dilated = F.max_pool2d(mask_f, kernel_size=11, stride=1, padding=5)
+                    dilated = F.max_pool2d(
+                        mask_f,
+                        kernel_size=RING_KERNEL_SIZE,
+                        stride=1,
+                        padding=RING_PADDING,
+                    )
                     ring_weight = (dilated - mask_f).clamp(min=0.0, max=1.0) * inpaint_weight
                 if isinstance(model_options, dict):
                     trace = model_options.get("lanpaint_semantic_trace")
 
+                if callable(distance_fn):
+                    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:
+                            # Use 3-arg positional: fn(prev, cur, ctx)
+                            dist_wrapper = lambda p, c, ctx: distance_fn(p, c, ctx)
+                        elif has_ctx_param or has_var_kw:
+                            # Use keyword: fn(prev, cur, ctx=ctx)
+                            dist_wrapper = lambda p, c, ctx: distance_fn(p, c, ctx=ctx)
+                        else:
+                            # Default 2-arg: fn(cur, prev)
+                            dist_wrapper = lambda p, c, ctx: distance_fn(c, p)
+                    except (ValueError, TypeError):
+                        # Fallback for built-ins or complex callables.
+                        # We try 3-args, then fall back to 2-args only if the TypeError
+                        # looks like an argument mismatch.
+                        def fallback_wrapper(p, c, ctx):
+                            try:
+                                return distance_fn(p, c, ctx)
+                            except TypeError as e:
+                                msg = str(e)
+                                if (
+                                    "positional argument" in msg
+                                    or "positional arguments" in msg
+                                    or "required positional" in msg
+                                    or ("takes" in msg and "given" in msg)
+                                ):
+                                    return distance_fn(c, p)
+                                raise
+
+                        dist_wrapper = fallback_wrapper
+
         x0_anchor = None
+
+        # Pre-fetch trace keys to avoid repeated dict lookups
+        bench_case_id = None
+        bench_outer_step = None
+        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")
+
         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 )
 
@@ -139,21 +209,20 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
                 x0_prev = None
                 x0_cur = None
 
-                if callable(distance_fn):
-                    try:
-                        dist = distance_fn(x_t_prev, x_t, ctx)
-                    except TypeError:
-                        dist = distance_fn(x_t, x_t_prev)
+                if dist_wrapper is not None:
+                    dist = dist_wrapper(x_t_prev, x_t, ctx)
                     custom_dist = dist is not None
 
                 if dist is None:
+                    # 'inpaint_weight' is guaranteed to be set when enabled_early_stop is True.
+                    inpaint = inpaint_weight
+
                     if isinstance(prev_args, tuple) and len(prev_args) >= 3:
                         x0_prev = prev_args[2]
                     if isinstance(args, tuple) and len(args) >= 3:
                         x0_cur = args[2]
 
                     if x0_prev is not None and x0_cur is not None:
-                        inpaint = inpaint_weight if inpaint_weight is not None else (1 - latent_mask).to(dtype=torch.float32)
                         diff_sq = (x0_cur.to(dtype=torch.float32) - x0_prev.to(dtype=torch.float32)) ** 2
                         denom = torch.sum(inpaint) + 1e-12
                         dist_inpaint = (torch.sum(diff_sq * inpaint) / denom).item()
@@ -163,7 +232,6 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
                             dist_ring = (torch.sum(diff_sq * ring_weight) / ring_denom).item()
                             dist = max(float(dist_inpaint), float(dist_ring))
                     else:
-                        inpaint = inpaint_weight if inpaint_weight is not None else (1 - latent_mask).to(dtype=torch.float32)
                         diff_sq = (x_t.to(dtype=torch.float32) - x_t_before.to(dtype=torch.float32)) ** 2
                         denom = torch.sum(inpaint) + 1e-12
                         dist = (torch.sum(diff_sq * inpaint) / denom).item()
@@ -178,7 +246,7 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
                         if x0_anchor is None:
                             x0_anchor = x0_cur.detach()
                         else:
-                            inpaint = inpaint_weight if inpaint_weight is not None else (1 - latent_mask).to(dtype=torch.float32)
+                            inpaint = inpaint_weight
                             diff_sq = (x0_cur.to(dtype=torch.float32) - x0_anchor.to(dtype=torch.float32)) ** 2
                             denom = torch.sum(inpaint) + 1e-12
                             drift_inpaint = (torch.sum(diff_sq * inpaint) / denom).item()
@@ -201,9 +269,9 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
                 if isinstance(trace, list):
                     trace.append(
                         {
-                            "case_id": model_options.get("bench_case_id") if isinstance(model_options, dict) else None,
-                            "outer_step": model_options.get("bench_outer_step") if isinstance(model_options, dict) else None,
-                            "bench_timestep": model_options.get("bench_timestep") if isinstance(model_options, dict) else None,
+                            "case_id": bench_case_id,
+                            "outer_step": bench_outer_step,
+                            "bench_timestep": bench_timestep,
                             "inner_step": i + 1,
                             "dist": float(dist),
                             "dist_inpaint": None if dist_inpaint is None else float(dist_inpaint),
diff --git a/tests/benchmark_lpips.py b/tests/benchmark_lpips.py
index 75fc297..a9f7f70 100644
--- a/tests/benchmark_lpips.py
+++ b/tests/benchmark_lpips.py
@@ -9,6 +9,7 @@
 import json
 import sys
 import time
+import warnings
 from pathlib import Path
 
 import torch
@@ -87,8 +88,8 @@ def _load_images(
     if shuffle:
         try:
             ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer)
-        except Exception:
-            pass
+        except Exception as e:
+            warnings.warn(f"Failed to shuffle streaming dataset; continuing without shuffle: {e}")
 
     images: list[torch.Tensor] = []
     for idx, item in enumerate(ds):
diff --git a/tests/compare_benchmarks.py b/tests/compare_benchmarks.py
index 8861233..b417747 100644
--- a/tests/compare_benchmarks.py
+++ b/tests/compare_benchmarks.py
@@ -4,6 +4,7 @@
 import json
 import math
 import os
+import warnings
 from pathlib import Path
 
 
@@ -100,7 +101,8 @@ def main() -> None:
                 try:
                     b_val = float(b_row[args.case_metric])
                     c_val = float(c_row[args.case_metric])
-                except Exception:
+                except Exception as e:
+                    warnings.warn(f"Skipping case {case_id} due to error: {e}")
                     continue
                 b_vals.append(b_val)
                 c_vals.append(c_val)
diff --git a/tests/render_benchmark_gallery.py b/tests/render_benchmark_gallery.py
index 8671a60..0875690 100644
--- a/tests/render_benchmark_gallery.py
+++ b/tests/render_benchmark_gallery.py
@@ -1,5 +1,6 @@
 import argparse
 import json
+import warnings
 from pathlib import Path
 
 import numpy as np
@@ -80,7 +81,8 @@ def main() -> None:
         try:
             b_val = float(b_cases[case_id][args.metric])
             c_val = float(c_cases[case_id][args.metric])
-        except Exception:
+        except Exception as e:
+            warnings.warn(f"Skipping case {case_id} due to error: {e}")
             continue
         deltas.append((c_val - b_val, case_id))
 
diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
index 88451c6..e5c9001 100644
--- a/tests/test_lanpaint_semantic_stop.py
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -84,6 +84,125 @@ def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
     assert calls["distance"] == 3
 
 
+def test_semantic_stop_distance_fn_two_arg_order_is_cur_prev() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    calls = {"langevin": 0, "distance": 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
+        return x_t + 1.0, args
+
+    def distance_fn(cur_x, prev_x):  # type: ignore[no-untyped-def]
+        assert float(torch.mean(cur_x).item()) > float(torch.mean(prev_x).item())
+        calls["distance"] += 1
+        return 0.0
+
+    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "min_steps": 3,
+            "patience": 1,
+            "distance_fn": distance_fn,
+        }
+    }
+
+    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["distance"] == 3
+
+
+def test_semantic_stop_distance_fn_accepts_third_param_not_named_ctx() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    calls = {"langevin": 0, "distance": 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
+        return x_t, args
+
+    def distance_fn(prev_x, cur_x, meta):  # type: ignore[no-untyped-def]
+        assert isinstance(meta, dict)
+        assert meta["step"] == calls["distance"]
+        calls["distance"] += 1
+        return 0.0
+
+    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "min_steps": 3,
+            "patience": 1,
+            "distance_fn": distance_fn,
+        }
+    }
+
+    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["distance"] == 3
+
+
+def test_semantic_stop_distance_fn_accepts_kw_only_ctx() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    calls = {"langevin": 0, "distance": 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
+        return x_t, args
+
+    def distance_fn(prev_x, cur_x, *, ctx):  # type: ignore[no-untyped-def]
+        assert isinstance(ctx, dict)
+        assert ctx["step"] == calls["distance"]
+        calls["distance"] += 1
+        return 0.0
+
+    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "min_steps": 3,
+            "patience": 1,
+            "distance_fn": distance_fn,
+        }
+    }
+
+    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["distance"] == 3
+
+
 def test_default_semantic_stop_triggers_at_min_steps_without_custom_distance_fn() -> None:
     engine = LanPaintEngine(
         _DummyModel(),
@@ -433,3 +552,64 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
     # If ring changes were ignored, inpaint-average MSE would fall below the threshold
     # and the sampler would stop early. This test ensures we keep running the model
     # while the boundary is still changing.
+
+
+def test_semantic_stop_distance_fn_dispatch_variants() -> None:
+    # 1. Standard 3-arg
+    def dist3(p, c, ctx):
+        dist3.called = True
+        return 0.0
+
+    dist3.called = False
+
+    # 2. 3-arg with different name
+    def dist3_alt(p, c, context):
+        dist3_alt.called = True
+        return 0.0
+
+    dist3_alt.called = False
+
+    # 3. 2-arg
+    def dist2(p, c):
+        dist2.called = True
+        return 0.0
+
+    dist2.called = False
+
+    # 4. Keyword-only ctx
+    def dist_kw(p, c, *, ctx):
+        dist_kw.called = True
+        return 0.0
+
+    dist_kw.called = False
+
+    # 5. Varargs
+    def dist_varargs(p, c, *args):
+        dist_varargs.called = True
+        return 0.0
+
+    dist_varargs.called = False
+
+    # 6. Varkwargs
+    def dist_varkw(p, c, **kwargs):
+        dist_varkw.called = True
+        assert "ctx" in kwargs
+        return 0.0
+
+    dist_varkw.called = False
+
+    variants = [dist3, dist3_alt, dist2, dist_kw, dist_varargs, dist_varkw]
+
+    for fn in variants:
+        engine = LanPaintEngine(_DummyModel(), NSteps=2, Friction=15.0, Lambda=1.0, Beta=1.0, StepSize=0.2)
+        model_options = {
+            "lanpaint_semantic_stop": {
+                "threshold": 1e-6,
+                "min_steps": 1,
+                "patience": 1,
+                "distance_fn": fn,
+            }
+        }
+        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=2)
+        assert fn.called, f"Function {fn.__name__} was not called"

From c5fe6351ea8df3703cf036847406d7281d9ea31d Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 19 Jan 2026 20:16:58 +0800
Subject: [PATCH 15/33] Harden semantic-stop distance hook

---
 src/LanPaint/lanpaint.py             |  24 ++++---
 tests/test_lanpaint_semantic_stop.py | 103 +++++++++++++++++++++++++++
 2 files changed, 116 insertions(+), 11 deletions(-)

diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index db4e265..0275ec8 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -1,4 +1,5 @@
 import torch
+from torch.nn import functional as F
 from .utils import StochasticHarmonicOscillator
 from functools import partial
 import inspect
@@ -72,6 +73,7 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
         patience = int(self.early_stop_patience)
         min_steps = 1
         distance_fn = self.early_stop_hook
+        # 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))
@@ -113,8 +115,6 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
 
                 inpaint_weight = (1 - latent_mask).to(dtype=torch.float32)
                 if latent_mask.dim() == 4:
-                    from torch.nn import functional as F
-
                     mask_f = latent_mask.to(dtype=torch.float32)
                     dilated = F.max_pool2d(
                         mask_f,
@@ -158,15 +158,10 @@ def fallback_wrapper(p, c, ctx):
                             try:
                                 return distance_fn(p, c, ctx)
                             except TypeError as e:
-                                msg = str(e)
-                                if (
-                                    "positional argument" in msg
-                                    or "positional arguments" in msg
-                                    or "required positional" in msg
-                                    or ("takes" in msg and "given" in msg)
-                                ):
-                                    return distance_fn(c, p)
-                                raise
+                                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)
 
                         dist_wrapper = fallback_wrapper
 
@@ -211,6 +206,13 @@ def fallback_wrapper(p, c, ctx):
 
                 if dist_wrapper is not None:
                     dist = dist_wrapper(x_t_prev, x_t, 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:
diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
index e5c9001..4f8dc2b 100644
--- a/tests/test_lanpaint_semantic_stop.py
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -203,6 +203,109 @@ def distance_fn(prev_x, cur_x, *, ctx):  # type: ignore[no-untyped-def]
     assert calls["distance"] == 3
 
 
+def test_semantic_stop_distance_fn_accepts_scalar_tensor_return() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    calls = {"langevin": 0, "distance": 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
+        return x_t, args
+
+    def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
+        calls["distance"] += 1
+        return torch.tensor(0.0)
+
+    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "min_steps": 3,
+            "patience": 1,
+            "distance_fn": distance_fn,
+        }
+    }
+
+    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["distance"] == 3
+
+
+def test_semantic_stop_distance_fn_rejects_multi_element_tensor_return() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
+        return torch.tensor([0.0, 0.0])
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "min_steps": 3,
+            "patience": 1,
+            "distance_fn": distance_fn,
+        }
+    }
+
+    x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
+
+    try:
+        engine(x, latent_image, noise, sigma, latent_mask, current_times, model_options=model_options, seed=0, n_steps=10)
+    except TypeError:
+        return
+    raise AssertionError("expected TypeError for multi-element tensor distance")
+
+
+def test_semantic_stop_fallback_wrapper_does_not_swallow_internal_typeerror() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    class DistanceFn:  # type: ignore[no-untyped-def]
+        __signature__ = "broken"
+
+        def __call__(self, prev_x, cur_x, ctx):
+            raise TypeError("internal")
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "min_steps": 3,
+            "patience": 1,
+            "distance_fn": DistanceFn(),
+        }
+    }
+
+    x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
+
+    try:
+        engine(x, latent_image, noise, sigma, latent_mask, current_times, model_options=model_options, seed=0, n_steps=10)
+    except TypeError:
+        return
+    raise AssertionError("expected TypeError from distance_fn to propagate")
+
+
 def test_default_semantic_stop_triggers_at_min_steps_without_custom_distance_fn() -> None:
     engine = LanPaintEngine(
         _DummyModel(),

From 7d46147887d28fa383dd88f942bca31add966294 Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 19 Jan 2026 20:21:32 +0800
Subject: [PATCH 16/33] Guard torch.nn.functional import

---
 src/LanPaint/lanpaint.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index 0275ec8..78b239c 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -1,5 +1,9 @@
 import torch
-from torch.nn import functional as F
+try:
+    from torch.nn import functional as F
+except Exception:
+    # Some environments (e.g. node graph validation) may not ship with a full torch package.
+    F = None  # type: ignore[assignment]
 from .utils import StochasticHarmonicOscillator
 from functools import partial
 import inspect
@@ -115,8 +119,12 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
 
                 inpaint_weight = (1 - latent_mask).to(dtype=torch.float32)
                 if latent_mask.dim() == 4:
+                    F_local = F
+                    if F_local is None:
+                        from torch.nn import functional as F_local
+
                     mask_f = latent_mask.to(dtype=torch.float32)
-                    dilated = F.max_pool2d(
+                    dilated = F_local.max_pool2d(
                         mask_f,
                         kernel_size=RING_KERNEL_SIZE,
                         stride=1,

From cbc9f3abfd7f704d1adbcba3a4e0d97d9708565f Mon Sep 17 00:00:00 2001
From: scraed <czhengac@connect.ust.hk>
Date: Thu, 22 Jan 2026 17:48:52 +0800
Subject: [PATCH 17/33] Refactor early-stop logic into LanPaintEarlyStopper
 class

Moved the early-stop logic from lanpaint.py into a new src/LanPaint/earlystop.py module as the LanPaintEarlyStopper class. This improves code organization and maintainability by encapsulating early-stop behavior, reducing complexity in the main LanPaint class.
---
 src/LanPaint/earlystop.py | 322 ++++++++++++++++++++++++++++++++++++++
 src/LanPaint/lanpaint.py  | 262 +++----------------------------
 2 files changed, 344 insertions(+), 240 deletions(-)
 create mode 100644 src/LanPaint/earlystop.py

diff --git a/src/LanPaint/earlystop.py b/src/LanPaint/earlystop.py
new file mode 100644
index 0000000..e3d733f
--- /dev/null
+++ b/src/LanPaint/earlystop.py
@@ -0,0 +1,322 @@
+import inspect
+from typing import Any, Callable, Optional
+
+import torch
+try:
+    from torch.nn import functional as F
+except Exception:
+    # Some environments (e.g. node graph validation) may not ship with a full torch package.
+    F = None  # type: ignore[assignment]
+
+# Early-stop constants
+STOP_THRESHOLD_MIN_ABT = 0.15
+STOP_THRESHOLD_MAX_ABT = 0.999
+PATIENCE_BOOST_LOW_ABT = 0.5
+PATIENCE_BOOST_HIGH_ABT = 0.9
+THRESHOLD_SCALE_MIN = 0.1
+RING_KERNEL_SIZE = 11
+RING_PADDING = 5
+
+
+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 = None
+        if isinstance(model_options, dict):
+            semantic_stop = model_options.get("lanpaint_semantic_stop")
+
+        threshold = float(default_threshold)
+        patience = int(default_patience)
+        min_steps = 1
+        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))
+            min_steps = int(semantic_stop.get("min_steps", min_steps))
+            distance_fn = semantic_stop.get("distance_fn", distance_fn)
+
+        enabled_early_stop = (threshold > 0.0) and (patience > 0)
+        min_steps = max(1, min_steps)
+        patience = max(1, patience)
+        patience_eff = patience
+        threshold_eff = threshold
+        inpaint_weight = None
+        ring_weight = None
+        trace = None
+        abt_val = None
+
+        if enabled_early_stop:
+            try:
+                abt_val = float(torch.mean(abt).item())
+            except Exception:
+                abt_val = 0.0
+
+            # Skip semantic early-stop in extremely noisy steps (low abt) and
+            # at the extreme tail where even tiny changes can matter.
+            if abt_val < STOP_THRESHOLD_MIN_ABT or abt_val > STOP_THRESHOLD_MAX_ABT:
+                enabled_early_stop = False
+            else:
+                # More noise -> require more consecutive stable steps before stopping.
+                # This keeps early-stop conservative in mid-noise outer steps.
+                patience_eff = patience + 1
+                if abt_val < PATIENCE_BOOST_LOW_ABT:
+                    patience_eff += 1
+                if abt_val > PATIENCE_BOOST_HIGH_ABT:
+                    patience_eff += 1
+                threshold_scale = max(THRESHOLD_SCALE_MIN, (1.0 - abt_val) ** 0.5)
+                threshold_eff = threshold * threshold_scale
+
+                inpaint_weight = (1 - latent_mask).to(dtype=torch.float32)
+                if latent_mask.dim() == 4:
+                    F_local = F
+                    if F_local is None:
+                        from torch.nn import functional as F_local
+
+                    mask_f = latent_mask.to(dtype=torch.float32)
+                    dilated = F_local.max_pool2d(
+                        mask_f,
+                        kernel_size=RING_KERNEL_SIZE,
+                        stride=1,
+                        padding=RING_PADDING,
+                    )
+                    ring_weight = (dilated - mask_f).clamp(min=0.0, max=1.0) * 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 = None
+        bench_outer_step = None
+        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,
+            min_steps=min_steps,
+            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,
+        min_steps: 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.min_steps = int(min_steps)
+
+        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]]):
+        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 = None
+        dist_ring = None
+        dist_drift = None
+        x0_prev = None
+        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:
+            if isinstance(prev_args, tuple) and len(prev_args) >= 3:
+                x0_prev = prev_args[2]
+            if isinstance(args, tuple) and len(args) >= 3:
+                x0_cur = args[2]
+
+            if x0_prev is not None and x0_cur is not None:
+                diff_sq = (x0_cur.to(dtype=torch.float32) - x0_prev.to(dtype=torch.float32)) ** 2
+                denom = torch.sum(inpaint) + 1e-12
+                dist_inpaint = (torch.sum(diff_sq * inpaint) / denom).item()
+                dist = float(dist_inpaint)
+                if self.ring_weight is not None:
+                    ring_denom = torch.sum(self.ring_weight) + 1e-12
+                    dist_ring = (torch.sum(diff_sq * self.ring_weight) / ring_denom).item()
+                    dist = max(float(dist_inpaint), float(dist_ring))
+            else:
+                diff_sq = (x_t_after.to(dtype=torch.float32) - x_t_before.to(dtype=torch.float32)) ** 2
+                denom = torch.sum(inpaint) + 1e-12
+                dist = (torch.sum(diff_sq * inpaint) / denom).item()
+                dist_inpaint = dist
+
+        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 float(dist) <= threshold_used:
+                if self.x0_anchor is None:
+                    self.x0_anchor = x0_cur.detach()
+                else:
+                    diff_sq = (x0_cur.to(dtype=torch.float32) - self.x0_anchor.to(dtype=torch.float32)) ** 2
+                    denom = torch.sum(inpaint) + 1e-12
+                    drift_inpaint = (torch.sum(diff_sq * inpaint) / denom).item()
+                    dist_drift = float(drift_inpaint)
+                    if self.ring_weight is not None:
+                        ring_denom = torch.sum(self.ring_weight) + 1e-12
+                        drift_ring = (torch.sum(diff_sq * self.ring_weight) / ring_denom).item()
+                        dist_drift = max(float(drift_inpaint), float(drift_ring))
+                    dist = max(float(dist), float(dist_drift))
+            else:
+                self.x0_anchor = None
+
+        if float(dist) <= threshold_used:
+            self.patience_counter += 1
+        else:
+            self.patience_counter = 0
+            self.x0_anchor = None
+
+        should_stop = (i + 1) >= self.min_steps and 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": float(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),
+                    "min_steps": int(self.min_steps),
+                    "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)
+
diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index 78b239c..37e4500 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -1,21 +1,7 @@
 import torch
-try:
-    from torch.nn import functional as F
-except Exception:
-    # Some environments (e.g. node graph validation) may not ship with a full torch package.
-    F = None  # type: ignore[assignment]
 from .utils import StochasticHarmonicOscillator
 from functools import partial
-import inspect
-
-# Early-stop constants
-STOP_THRESHOLD_MIN_ABT = 0.15
-STOP_THRESHOLD_MAX_ABT = 0.999
-PATIENCE_BOOST_LOW_ABT = 0.5
-PATIENCE_BOOST_HIGH_ABT = 0.9
-THRESHOLD_SCALE_MIN = 0.1
-RING_KERNEL_SIZE = 11
-RING_PADDING = 5
+from .earlystop import LanPaintEarlyStopper
 
 class LanPaint():
     def __init__(self, Model, NSteps, Friction, Lambda, Beta, StepSize, IS_FLUX = False, IS_FLOW = False, EarlyStopThreshold = 0.0, EarlyStopPatience = 1, EarlyStopHook = None):
@@ -69,131 +55,25 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
         ############ 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
-        semantic_stop = None
-        if isinstance(model_options, dict):
-            semantic_stop = model_options.get("lanpaint_semantic_stop")
-
-        threshold = float(self.early_stop_threshold)
-        patience = int(self.early_stop_patience)
-        min_steps = 1
-        distance_fn = self.early_stop_hook
-        # 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))
-            min_steps = int(semantic_stop.get("min_steps", min_steps))
-            distance_fn = semantic_stop.get("distance_fn", distance_fn)
-
-        enabled_early_stop = (threshold > 0.0) and (patience > 0)
-        min_steps = max(1, min_steps)
-        patience = max(1, patience)
-        patience_counter = 0
-        patience_eff = patience
-        threshold_eff = threshold
-        inpaint_weight = None
-        ring_weight = None
-        trace = None
-        dist_wrapper = None
-
-        if enabled_early_stop:
-            try:
-                abt_val = float(torch.mean(abt).item())
-            except Exception:
-                abt_val = 0.0
-
-            # Skip semantic early-stop in extremely noisy steps (low abt) and
-            # at the extreme tail where even tiny changes can matter.
-            if abt_val < STOP_THRESHOLD_MIN_ABT or abt_val > STOP_THRESHOLD_MAX_ABT:
-                enabled_early_stop = False
-            else:
-                # More noise -> require more consecutive stable steps before stopping.
-                # This keeps early-stop conservative in mid-noise outer steps.
-                patience_eff = patience + 1
-                if abt_val < PATIENCE_BOOST_LOW_ABT:
-                    patience_eff += 1
-                if abt_val > PATIENCE_BOOST_HIGH_ABT:
-                    patience_eff += 1
-                threshold_scale = max(THRESHOLD_SCALE_MIN, (1.0 - abt_val) ** 0.5)
-                threshold_eff = threshold * threshold_scale
-
-                inpaint_weight = (1 - latent_mask).to(dtype=torch.float32)
-                if latent_mask.dim() == 4:
-                    F_local = F
-                    if F_local is None:
-                        from torch.nn import functional as F_local
-
-                    mask_f = latent_mask.to(dtype=torch.float32)
-                    dilated = F_local.max_pool2d(
-                        mask_f,
-                        kernel_size=RING_KERNEL_SIZE,
-                        stride=1,
-                        padding=RING_PADDING,
-                    )
-                    ring_weight = (dilated - mask_f).clamp(min=0.0, max=1.0) * inpaint_weight
-                if isinstance(model_options, dict):
-                    trace = model_options.get("lanpaint_semantic_trace")
-
-                if callable(distance_fn):
-                    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:
-                            # Use 3-arg positional: fn(prev, cur, ctx)
-                            dist_wrapper = lambda p, c, ctx: distance_fn(p, c, ctx)
-                        elif has_ctx_param or has_var_kw:
-                            # Use keyword: fn(prev, cur, ctx=ctx)
-                            dist_wrapper = lambda p, c, ctx: distance_fn(p, c, ctx=ctx)
-                        else:
-                            # Default 2-arg: fn(cur, prev)
-                            dist_wrapper = lambda p, c, ctx: distance_fn(c, p)
-                    except (ValueError, TypeError):
-                        # Fallback for built-ins or complex callables.
-                        # We try 3-args, then fall back to 2-args only if the TypeError
-                        # looks like an argument mismatch.
-                        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)
-
-                        dist_wrapper = fallback_wrapper
-
-        x0_anchor = None
-
-        # Pre-fetch trace keys to avoid repeated dict lookups
-        bench_case_id = None
-        bench_outer_step = None
-        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")
+        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 enabled_early_stop and callable(distance_fn) else None
-            x_t_before = x_t if enabled_early_stop else None
+            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 enabled_early_stop and x_t_before is not None:
+            if stopper is not None and x_t_before is not None:
                 ctx = {
                     "step": i,
                     "steps_done": i + 1,
@@ -203,102 +83,16 @@ def fallback_wrapper(p, c, ctx):
                     "current_times": current_times,
                     "seed": seed,
                 }
-
-                dist = None
-                custom_dist = False
-                dist_inpaint = None
-                dist_ring = None
-                dist_drift = None
-                x0_prev = None
-                x0_cur = None
-
-                if dist_wrapper is not None:
-                    dist = dist_wrapper(x_t_prev, x_t, 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:
-                    # 'inpaint_weight' is guaranteed to be set when enabled_early_stop is True.
-                    inpaint = inpaint_weight
-
-                    if isinstance(prev_args, tuple) and len(prev_args) >= 3:
-                        x0_prev = prev_args[2]
-                    if isinstance(args, tuple) and len(args) >= 3:
-                        x0_cur = args[2]
-
-                    if x0_prev is not None and x0_cur is not None:
-                        diff_sq = (x0_cur.to(dtype=torch.float32) - x0_prev.to(dtype=torch.float32)) ** 2
-                        denom = torch.sum(inpaint) + 1e-12
-                        dist_inpaint = (torch.sum(diff_sq * inpaint) / denom).item()
-                        dist = float(dist_inpaint)
-                        if ring_weight is not None:
-                            ring_denom = torch.sum(ring_weight) + 1e-12
-                            dist_ring = (torch.sum(diff_sq * ring_weight) / ring_denom).item()
-                            dist = max(float(dist_inpaint), float(dist_ring))
-                    else:
-                        diff_sq = (x_t.to(dtype=torch.float32) - x_t_before.to(dtype=torch.float32)) ** 2
-                        denom = torch.sum(inpaint) + 1e-12
-                        dist = (torch.sum(diff_sq * inpaint) / denom).item()
-                        dist_inpaint = dist
-
-                threshold_used = threshold if custom_dist else threshold_eff
-
-                # Extra guard: even if per-step changes are tiny, avoid stopping if the
-                # x0 estimate keeps drifting over several steps.
-                if x0_cur is not None and not custom_dist:
-                    if float(dist) <= threshold_used:
-                        if x0_anchor is None:
-                            x0_anchor = x0_cur.detach()
-                        else:
-                            inpaint = inpaint_weight
-                            diff_sq = (x0_cur.to(dtype=torch.float32) - x0_anchor.to(dtype=torch.float32)) ** 2
-                            denom = torch.sum(inpaint) + 1e-12
-                            drift_inpaint = (torch.sum(diff_sq * inpaint) / denom).item()
-                            dist_drift = float(drift_inpaint)
-                            if ring_weight is not None:
-                                ring_denom = torch.sum(ring_weight) + 1e-12
-                                drift_ring = (torch.sum(diff_sq * ring_weight) / ring_denom).item()
-                                dist_drift = max(float(drift_inpaint), float(drift_ring))
-                            dist = max(float(dist), float(dist_drift))
-                    else:
-                        x0_anchor = None
-
-                if float(dist) <= threshold_used:
-                    patience_counter += 1
-                else:
-                    patience_counter = 0
-                    x0_anchor = None
-
-                should_stop = (i + 1) >= min_steps and patience_counter >= patience_eff
-                if isinstance(trace, list):
-                    trace.append(
-                        {
-                            "case_id": bench_case_id,
-                            "outer_step": bench_outer_step,
-                            "bench_timestep": bench_timestep,
-                            "inner_step": i + 1,
-                            "dist": float(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(threshold_eff),
-                            "patience_counter": int(patience_counter),
-                            "patience_eff": int(patience_eff),
-                            "min_steps": int(min_steps),
-                            "abt": float(abt_val),
-                            "custom_dist": bool(custom_dist),
-                            "stopped": bool(should_stop),
-                        }
-                    )
-
-                if should_stop:
+                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:
@@ -358,18 +152,6 @@ def advance_time(x_t, v, dt, Gamma, A, C, D):
             v = v.to(dtype)
             return x_t, v
 
-        if score is None:
-            if not (isinstance(args, tuple) and len(args) >= 3):
-                return x_t, args
-            v = args[0]
-            C = args[1]
-            x0 = args[2]
-            x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D)
-            C_new = (abt**0.5 * x0  - x_t )/ (1-abt) + A * x_t
-            v = v + Gamma**0.5 * ( C_new - C) *dt
-            x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C_new, D)
-            return x_t, (v, C_new, x0)
-
         def Coef_C(x_t):
             x0 = self.x0_evalutation(x_t, score, sigma, args)
             C = (abt**0.5 * x0  - x_t )/ (1-abt) + A * x_t

From 1b87239553c4a62d3d46b66e8d81aefdcb5e2e51 Mon Sep 17 00:00:00 2001
From: scraed <czhengac@connect.ust.hk>
Date: Thu, 22 Jan 2026 17:51:05 +0800
Subject: [PATCH 18/33] Move Coef_C definition above advance_time

Reordered the Coef_C function definition to appear before advance_time for improved code organization and readability.
---
 src/LanPaint/lanpaint.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index 37e4500..0bc99d3 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -143,6 +143,10 @@ def langevin_dynamics(self, x_t, score, mask, step_size, current_times, sigma_x=
         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)
+            C = (abt**0.5 * x0  - x_t )/ (1-abt) + A * x_t
+            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):
@@ -152,11 +156,6 @@ def advance_time(x_t, v, dt, Gamma, A, C, D):
             v = v.to(dtype)
             return x_t, v
 
-        def Coef_C(x_t):
-            x0 = self.x0_evalutation(x_t, score, sigma, args)
-            C = (abt**0.5 * x0  - x_t )/ (1-abt) + A * x_t
-            return C, x0
-
         def advance_time_overdamped(x_t, dt, A, C, D):
             """
             Overdamped (Gamma -> infinity) limit:

From 34ded92d8d976e963f5269918a468181ac3d9593 Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Sat, 24 Jan 2026 10:16:13 +0800
Subject: [PATCH 19/33] Simplify semantic early-stop knobs

---
 .github/workflows/benchmark.yml      |   5 -
 src/LanPaint/earlystop.py            |  93 ++++----
 tests/benchmark_lpips.py             |   5 -
 tests/test_lanpaint_semantic_stop.py | 343 +++++++++++++++++++++------
 4 files changed, 316 insertions(+), 130 deletions(-)

diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml
index dfa0098..df21edb 100644
--- a/.github/workflows/benchmark.yml
+++ b/.github/workflows/benchmark.yml
@@ -35,10 +35,6 @@ on:
         description: "Semantic early-stop patience (head_on run)"
         required: true
         default: "2"
-      semantic_min_steps:
-        description: "Semantic early-stop minimum steps (head_on run)"
-        required: true
-        default: "2"
       min_seconds:
         description: "Fail if a benchmark runtime is below this (0 to disable)"
         required: true
@@ -191,7 +187,6 @@ jobs:
             --min-seconds "${{ inputs.min_seconds }}" \
             --semantic-threshold "${{ inputs.semantic_threshold }}" \
             --semantic-patience "${{ inputs.semantic_patience }}" \
-            --semantic-min-steps "${{ inputs.semantic_min_steps }}" \
             --save-images-dir head_on_images
 
       - name: Compare baseline vs head_off
diff --git a/src/LanPaint/earlystop.py b/src/LanPaint/earlystop.py
index e3d733f..d526433 100644
--- a/src/LanPaint/earlystop.py
+++ b/src/LanPaint/earlystop.py
@@ -2,20 +2,45 @@
 from typing import Any, Callable, Optional
 
 import torch
-try:
-    from torch.nn import functional as F
-except Exception:
-    # Some environments (e.g. node graph validation) may not ship with a full torch package.
-    F = None  # type: ignore[assignment]
-
-# Early-stop constants
-STOP_THRESHOLD_MIN_ABT = 0.15
-STOP_THRESHOLD_MAX_ABT = 0.999
-PATIENCE_BOOST_LOW_ABT = 0.5
-PATIENCE_BOOST_HIGH_ABT = 0.9
-THRESHOLD_SCALE_MIN = 0.1
-RING_KERNEL_SIZE = 11
-RING_PADDING = 5
+
+
+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
 
 
 class LanPaintEarlyStopper:
@@ -40,20 +65,17 @@ def from_options(
 
         threshold = float(default_threshold)
         patience = int(default_patience)
-        min_steps = 1
         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))
-            min_steps = int(semantic_stop.get("min_steps", min_steps))
             distance_fn = semantic_stop.get("distance_fn", distance_fn)
 
         enabled_early_stop = (threshold > 0.0) and (patience > 0)
-        min_steps = max(1, min_steps)
         patience = max(1, patience)
-        patience_eff = patience
+        patience_eff = patience + 1
         threshold_eff = threshold
         inpaint_weight = None
         ring_weight = None
@@ -66,35 +88,12 @@ def from_options(
             except Exception:
                 abt_val = 0.0
 
-            # Skip semantic early-stop in extremely noisy steps (low abt) and
-            # at the extreme tail where even tiny changes can matter.
-            if abt_val < STOP_THRESHOLD_MIN_ABT or abt_val > STOP_THRESHOLD_MAX_ABT:
+            threshold_eff = threshold * _abt_scale(abt_val)
+            if threshold_eff <= 0.0:
                 enabled_early_stop = False
             else:
-                # More noise -> require more consecutive stable steps before stopping.
-                # This keeps early-stop conservative in mid-noise outer steps.
-                patience_eff = patience + 1
-                if abt_val < PATIENCE_BOOST_LOW_ABT:
-                    patience_eff += 1
-                if abt_val > PATIENCE_BOOST_HIGH_ABT:
-                    patience_eff += 1
-                threshold_scale = max(THRESHOLD_SCALE_MIN, (1.0 - abt_val) ** 0.5)
-                threshold_eff = threshold * threshold_scale
-
                 inpaint_weight = (1 - latent_mask).to(dtype=torch.float32)
-                if latent_mask.dim() == 4:
-                    F_local = F
-                    if F_local is None:
-                        from torch.nn import functional as F_local
-
-                    mask_f = latent_mask.to(dtype=torch.float32)
-                    dilated = F_local.max_pool2d(
-                        mask_f,
-                        kernel_size=RING_KERNEL_SIZE,
-                        stride=1,
-                        padding=RING_PADDING,
-                    )
-                    ring_weight = (dilated - mask_f).clamp(min=0.0, max=1.0) * inpaint_weight
+                ring_weight = _boundary_weight(latent_mask, inpaint_weight)
                 if isinstance(model_options, dict):
                     trace = model_options.get("lanpaint_semantic_trace")
 
@@ -115,7 +114,6 @@ def from_options(
             threshold=threshold,
             threshold_eff=threshold_eff,
             patience_eff=patience_eff,
-            min_steps=min_steps,
             inpaint_weight=inpaint_weight,
             ring_weight=ring_weight,
             distance_fn=distance_fn,
@@ -133,7 +131,6 @@ def __init__(
         threshold: float,
         threshold_eff: float,
         patience_eff: int,
-        min_steps: int,
         inpaint_weight: Optional[torch.Tensor],
         ring_weight: Optional[torch.Tensor],
         distance_fn: Optional[Callable[..., Any]] = None,
@@ -147,7 +144,6 @@ def __init__(
         self.threshold = float(threshold)
         self.threshold_eff = float(threshold_eff)
         self.patience_eff = int(patience_eff)
-        self.min_steps = int(min_steps)
 
         self.inpaint_weight = inpaint_weight
         self.ring_weight = ring_weight
@@ -294,7 +290,7 @@ def step(
             self.patience_counter = 0
             self.x0_anchor = None
 
-        should_stop = (i + 1) >= self.min_steps and self.patience_counter >= self.patience_eff
+        should_stop = self.patience_counter >= self.patience_eff
 
         if isinstance(self.trace, list):
             self.trace.append(
@@ -311,7 +307,6 @@ def step(
                     "threshold_eff": float(self.threshold_eff),
                     "patience_counter": int(self.patience_counter),
                     "patience_eff": int(self.patience_eff),
-                    "min_steps": int(self.min_steps),
                     "abt": None if self.abt_val is None else float(self.abt_val),
                     "custom_dist": bool(custom_dist),
                     "stopped": bool(should_stop),
diff --git a/tests/benchmark_lpips.py b/tests/benchmark_lpips.py
index a9f7f70..c85f9fe 100644
--- a/tests/benchmark_lpips.py
+++ b/tests/benchmark_lpips.py
@@ -525,7 +525,6 @@ def main() -> None:
         help="Enable LanPaint semantic early-stop if > 0 (guided_diffusion_e2e)",
     )
     parser.add_argument("--semantic-patience", type=int, default=1, help="Semantic early-stop patience (guided_diffusion_e2e)")
-    parser.add_argument("--semantic-min-steps", type=int, default=1, help="Semantic early-stop min_steps (guided_diffusion_e2e)")
     parser.add_argument("--report-earlystop", action="store_true", help="Include early-stop stats in output JSON (guided_diffusion_e2e)")
     parser.add_argument("--trace-inner", type=Path, default=None, help="If set, write per-inner-step semantic trace JSON (guided_diffusion_e2e)")
     parser.add_argument("--save-images-dir", type=Path, default=None, help="If set, save per-case images for manual review")
@@ -539,8 +538,6 @@ def main() -> None:
         raise SystemExit("--shuffle-buffer must be > 0")
     if args.semantic_patience <= 0:
         raise SystemExit("--semantic-patience must be > 0")
-    if args.semantic_min_steps <= 0:
-        raise SystemExit("--semantic-min-steps must be > 0")
     if args.outer_cutoff_steps < 0:
         raise SystemExit("--outer-cutoff-steps must be >= 0")
 
@@ -584,7 +581,6 @@ def main() -> None:
             semantic_stop = {
                 "threshold": float(args.semantic_threshold),
                 "patience": int(args.semantic_patience),
-                "min_steps": int(args.semantic_min_steps),
             }
 
         semantic_trace = [] if args.trace_inner is not None else None
@@ -639,7 +635,6 @@ def main() -> None:
                     "device": device_str,
                     "semantic_threshold": args.semantic_threshold,
                     "semantic_patience": args.semantic_patience,
-                    "semantic_min_steps": args.semantic_min_steps,
                     "outer_cutoff_steps": args.outer_cutoff_steps,
                     "report_earlystop": bool(args.report_earlystop),
                     "trace_inner": str(args.trace_inner) if args.trace_inner is not None else None,
diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
index 4f8dc2b..d087b0f 100644
--- a/tests/test_lanpaint_semantic_stop.py
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -1,5 +1,4 @@
 import torch
-import torch.nn.functional as F
 
 from src.LanPaint.lanpaint import LanPaint as LanPaintEngine
 
@@ -29,7 +28,7 @@ def _inputs():  # type: ignore[no-untyped-def]
     return x, latent_image, noise, sigma, latent_mask, current_times
 
 
-def test_semantic_stop_triggers_deterministically_at_min_steps() -> None:
+def test_semantic_stop_triggers_deterministically_at_patience() -> None:
     engine = LanPaintEngine(
         _DummyModel(),
         NSteps=10,
@@ -59,8 +58,7 @@ def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
     model_options = {
         "lanpaint_semantic_stop": {
             "threshold": 1e-6,
-            "min_steps": 3,
-            "patience": 1,
+            "patience": 2,
             "distance_fn": distance_fn,
         }
     }
@@ -84,6 +82,130 @@ def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
     assert calls["distance"] == 3
 
 
+def test_semantic_stop_ignores_min_steps_option() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    calls = {"langevin": 0, "distance": 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
+        return x_t, args
+
+    def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
+        calls["distance"] += 1
+        return 0.0
+
+    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "patience": 1,
+            "min_steps": 5,
+            "distance_fn": distance_fn,
+        }
+    }
+
+    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)
+
+    # 'min_steps' is ignored; stopping is controlled only by patience.
+    assert calls["langevin"] == 2
+    assert calls["distance"] == 2
+
+
+def test_semantic_stop_patience_is_not_abt_boosted() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    calls = {"langevin": 0, "distance": 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
+        return x_t, args
+
+    def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
+        calls["distance"] += 1
+        return 0.0
+
+    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "patience": 1,
+            "distance_fn": distance_fn,
+        }
+    }
+
+    x, latent_image, noise, sigma, latent_mask, _ = _inputs()
+
+    calls["langevin"] = 0
+    calls["distance"] = 0
+    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
+    engine(x, latent_image, noise, sigma, latent_mask, current_times, model_options=model_options, seed=0, n_steps=10)
+    assert calls["langevin"] == 2
+    assert calls["distance"] == 2
+
+    calls["langevin"] = 0
+    calls["distance"] = 0
+    current_times = (sigma, torch.tensor([0.99]), torch.tensor([0.0]))
+    engine(x, latent_image, noise, sigma, latent_mask, current_times, model_options=model_options, seed=0, n_steps=10)
+    assert calls["langevin"] == 2
+    assert calls["distance"] == 2
+
+
+def test_semantic_stop_is_not_disabled_below_old_min_abt_gate() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    calls = {"langevin": 0, "distance": 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
+        return x_t, args
+
+    def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
+        calls["distance"] += 1
+        return 0.0
+
+    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": 1e-6,
+            "patience": 1,
+            "distance_fn": distance_fn,
+        }
+    }
+
+    x, latent_image, noise, sigma, latent_mask, _ = _inputs()
+    current_times = (sigma, torch.tensor([0.1]), torch.tensor([0.0]))
+    engine(x, latent_image, noise, sigma, latent_mask, current_times, model_options=model_options, seed=0, n_steps=10)
+
+    assert calls["langevin"] == 2
+    assert calls["distance"] == 2
+
+
 def test_semantic_stop_distance_fn_two_arg_order_is_cur_prev() -> None:
     engine = LanPaintEngine(
         _DummyModel(),
@@ -110,8 +232,7 @@ def distance_fn(cur_x, prev_x):  # type: ignore[no-untyped-def]
     model_options = {
         "lanpaint_semantic_stop": {
             "threshold": 1e-6,
-            "min_steps": 3,
-            "patience": 1,
+            "patience": 2,
             "distance_fn": distance_fn,
         }
     }
@@ -150,8 +271,7 @@ def distance_fn(prev_x, cur_x, meta):  # type: ignore[no-untyped-def]
     model_options = {
         "lanpaint_semantic_stop": {
             "threshold": 1e-6,
-            "min_steps": 3,
-            "patience": 1,
+            "patience": 2,
             "distance_fn": distance_fn,
         }
     }
@@ -190,8 +310,7 @@ def distance_fn(prev_x, cur_x, *, ctx):  # type: ignore[no-untyped-def]
     model_options = {
         "lanpaint_semantic_stop": {
             "threshold": 1e-6,
-            "min_steps": 3,
-            "patience": 1,
+            "patience": 2,
             "distance_fn": distance_fn,
         }
     }
@@ -228,8 +347,7 @@ def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
     model_options = {
         "lanpaint_semantic_stop": {
             "threshold": 1e-6,
-            "min_steps": 3,
-            "patience": 1,
+            "patience": 2,
             "distance_fn": distance_fn,
         }
     }
@@ -257,8 +375,7 @@ def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
     model_options = {
         "lanpaint_semantic_stop": {
             "threshold": 1e-6,
-            "min_steps": 3,
-            "patience": 1,
+            "patience": 2,
             "distance_fn": distance_fn,
         }
     }
@@ -291,7 +408,6 @@ def __call__(self, prev_x, cur_x, ctx):
     model_options = {
         "lanpaint_semantic_stop": {
             "threshold": 1e-6,
-            "min_steps": 3,
             "patience": 1,
             "distance_fn": DistanceFn(),
         }
@@ -306,7 +422,7 @@ def __call__(self, prev_x, cur_x, ctx):
     raise AssertionError("expected TypeError from distance_fn to propagate")
 
 
-def test_default_semantic_stop_triggers_at_min_steps_without_custom_distance_fn() -> None:
+def test_default_semantic_stop_triggers_at_patience_without_custom_distance_fn() -> None:
     engine = LanPaintEngine(
         _DummyModel(),
         NSteps=10,
@@ -331,8 +447,7 @@ def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y
     model_options = {
         "lanpaint_semantic_stop": {
             "threshold": 1e-6,
-            "min_steps": 3,
-            "patience": 1,
+            "patience": 2,
         }
     }
 
@@ -344,7 +459,7 @@ def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y
     assert calls["without_score"] == 0
 
 
-def test_semantic_stop_is_disabled_for_very_late_steps() -> None:
+def test_semantic_stop_is_more_conservative_near_tail_abt() -> None:
     engine = LanPaintEngine(
         _DummyModel(),
         NSteps=10,
@@ -356,29 +471,13 @@ def test_semantic_stop_is_disabled_for_very_late_steps() -> None:
 
     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,
-            "min_steps": 2,
-            "patience": 1,
-        }
-    }
-
     x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
 
-    # Use a constant per-step x_t update (simulating Langevin noise) while keeping
-    # an x0-like tensor perfectly stable. Early-stop should ignore the x_t noise.
-    delta = 0.01
+    # Construct an x0 delta (MSE) that is below the mid-step effective threshold
+    # but above the near-tail effective threshold. This ensures early-stop is
+    # more conservative near the tail via threshold scaling.
+    delta_sq = 1e-7
+    delta = float(delta_sq**0.5)
 
     def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None):  # type: ignore[no-untyped-def]
         calls["langevin"] += 1
@@ -386,16 +485,18 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
             calls["without_score"] += 1
         else:
             calls["with_score"] += 1
-        prev_x0 = None if args is None else args[2]
-        x0 = x_t.detach().clone() * 0 if prev_x0 is None else prev_x0
-        return x_t + delta, (None, None, x0)
+
+        step = float(calls["langevin"])
+        x0 = torch.zeros_like(x_t) + (step * delta)
+
+        # Ensure the first iteration does not look stable due to x_t noise metric.
+        return x_t + (1.0 if step <= 1 else 0.0), (None, None, x0)
 
     engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
 
     model_options = {
         "lanpaint_semantic_stop": {
-            "threshold": 1e-12,
-            "min_steps": 2,
+            "threshold": 1e-6,
             "patience": 1,
         }
     }
@@ -432,8 +533,8 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         seed=0,
         n_steps=10,
     )
-    assert calls["langevin"] == 4
-    assert calls["with_score"] == 4
+    assert calls["langevin"] == 10
+    assert calls["with_score"] == 10
     assert calls["without_score"] == 0
 
 
@@ -462,7 +563,6 @@ def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y
     model_options = {
         "lanpaint_semantic_stop": {
             "threshold": 1e-6,
-            "min_steps": 2,
             "patience": 1,
         }
     }
@@ -512,17 +612,17 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
 
     engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
 
-    # Choose a threshold that would previously depend on how much of the tensor is masked.
+    # Choose a threshold that would diverge if the metric was accidentally area-scaled
+    # (sum instead of mean) over the inpaint region.
     model_options = {
         "lanpaint_semantic_stop": {
-            "threshold": (delta * delta) * 0.2,
-            "min_steps": 2,
+            "threshold": 0.005,
             "patience": 1,
         }
     }
 
     x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
-    current_times = (sigma, torch.tensor([0.0]), torch.tensor([0.0]))
+    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
 
     # Small hole (unknown region) -> mask mostly 1, with a small 0 patch.
     small_hole_mask = torch.ones_like(latent_mask)
@@ -542,8 +642,8 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         seed=0,
         n_steps=10,
     )
-    assert calls["langevin"] == 10
-    assert calls["with_score"] == 10
+    assert calls["langevin"] == 2
+    assert calls["with_score"] == 2
     assert calls["without_score"] == 0
 
     # Larger hole (unknown region).
@@ -564,12 +664,12 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         seed=0,
         n_steps=10,
     )
-    assert calls["langevin"] == 10
-    assert calls["with_score"] == 10
+    assert calls["langevin"] == 2
+    assert calls["with_score"] == 2
     assert calls["without_score"] == 0
 
 
-def test_semantic_stop_does_not_stop_while_ring_changes() -> None:
+def test_semantic_stop_does_not_stop_while_boundary_changes() -> None:
     engine = LanPaintEngine(
         _DummyModel(),
         NSteps=10,
@@ -594,19 +694,24 @@ def test_semantic_stop_does_not_stop_while_ring_changes() -> None:
     latent_mask[:, :, 12:52, 12:52] = 0.0
     inpaint_mask = 1.0 - latent_mask
 
-    # Replicate the "ring" construction in LanPaint: a dilated mask boundary
-    # intersected with the inpaint region.
-    mask_f = latent_mask.to(dtype=torch.float32)
-    dilated = F.max_pool2d(mask_f, kernel_size=11, stride=1, padding=5)
-    ring_mask = (dilated - mask_f).clamp(min=0.0, max=1.0) * inpaint_mask.to(dtype=torch.float32)
+    # Construct a 4-neighbor boundary: unknown pixels adjacent to known pixels.
+    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_mask = (inpaint_mask > 0.5) & neighbor_known
+    boundary_weight = boundary_mask.to(dtype=torch.float32)
 
-    ring_area = float(torch.sum(ring_mask).item())
+    boundary_area = float(torch.sum(boundary_weight).item())
     inpaint_area = float(torch.sum(inpaint_mask).item())
-    assert ring_area > 0.0
-    assert ring_area < inpaint_area
+    assert boundary_area > 0.0
+    assert boundary_area < inpaint_area
 
     threshold = 1e-6
-    ratio = inpaint_area / ring_area
+    ratio = inpaint_area / boundary_area
     delta_sq = threshold * ratio * 0.9
     delta = float(delta_sq**0.5)
 
@@ -619,18 +724,17 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         else:
             calls["with_score"] += 1
 
-        # Simulate a case where the interior is stable, but the boundary ring keeps changing.
+        # Simulate a case where the interior is stable, but the boundary keeps changing.
         step = float(calls["langevin"])
-        x0 = ring_mask.to(dtype=x_t.dtype) * (step * delta)
+        x0 = boundary_weight.to(dtype=x_t.dtype) * (step * delta)
         return x_t, (None, None, x0)
 
     engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
 
-    current_times = (sigma, torch.tensor([0.0]), torch.tensor([0.0]))
+    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
     model_options = {
         "lanpaint_semantic_stop": {
             "threshold": threshold,
-            "min_steps": 2,
             "patience": 1,
         }
     }
@@ -657,6 +761,104 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
     # while the boundary is still changing.
 
 
+def test_semantic_stop_boundary_guard_is_adjacency_only() -> None:
+    engine = LanPaintEngine(
+        _DummyModel(),
+        NSteps=10,
+        Friction=15.0,
+        Lambda=1.0,
+        Beta=1.0,
+        StepSize=0.2,
+    )
+
+    batch = 1
+    channels = 4
+    height = 64
+    width = 64
+
+    x = torch.zeros((batch, channels, height, width))
+    latent_image = torch.zeros_like(x)
+    noise = torch.ones_like(x)
+    sigma = torch.tensor([1.0])
+
+    # latent_mask uses LanPaint convention: 1=known region, 0=to-inpaint.
+    latent_mask = torch.ones_like(x)
+    h0, h1 = 12, 52
+    w0, w1 = 12, 52
+    latent_mask[:, :, h0:h1, w0:w1] = 0.0
+    inpaint_mask = 1.0 - latent_mask
+
+    # Pixels at distance 3 from the boundary (inside the old 5px ring, but not in the 1px adjacency boundary).
+    offset = 3
+    changed_mask = torch.zeros_like(inpaint_mask)
+    changed_mask[:, :, h0 + offset : h0 + offset + 1, w0 + offset : w1 - offset] = 1.0
+    changed_mask[:, :, h1 - offset - 1 : h1 - offset, w0 + offset : w1 - offset] = 1.0
+    changed_mask[:, :, h0 + offset + 1 : h1 - offset - 1, w0 + offset : w0 + offset + 1] = 1.0
+    changed_mask[:, :, h0 + offset + 1 : h1 - offset - 1, w1 - offset - 1 : w1 - offset] = 1.0
+
+    # Reference old "ring": all pixels within 5px of the boundary inside the hole.
+    ring_width = 5
+    ring5_mask = torch.zeros_like(inpaint_mask)
+    ring5_mask[:, :, h0 : h0 + ring_width, w0:w1] = 1.0
+    ring5_mask[:, :, h1 - ring_width : h1, w0:w1] = 1.0
+    ring5_mask[:, :, h0 + ring_width : h1 - ring_width, w0 : w0 + ring_width] = 1.0
+    ring5_mask[:, :, h0 + ring_width : h1 - ring_width, w1 - ring_width : w1] = 1.0
+
+    changed_area = float(torch.sum(changed_mask).item())
+    ring5_area = float(torch.sum(ring5_mask).item())
+    inpaint_area = float(torch.sum(inpaint_mask).item())
+    assert changed_area > 0.0
+    assert ring5_area > changed_area
+    assert inpaint_area > ring5_area
+
+    threshold = 1e-6
+    # Choose a delta so that the inpaint-average MSE falls below threshold, but the old 5px ring-average stays above.
+    delta_sq = threshold * (ring5_area / changed_area) * 1.1
+    assert (delta_sq * changed_area / inpaint_area) < threshold
+    assert (delta_sq * changed_area / ring5_area) > threshold
+    delta = float(delta_sq**0.5)
+
+    calls = {"langevin": 0, "with_score": 0, "without_score": 0}
+
+    def fake_langevin_delta(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
+
+        step = float(calls["langevin"])
+        x0 = changed_mask.to(dtype=x_t.dtype) * (step * delta)
+        return x_t, (None, None, x0)
+
+    engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
+
+    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": threshold,
+            "patience": 1,
+        }
+    }
+
+    engine(
+        x,
+        latent_image,
+        noise,
+        sigma,
+        latent_mask,
+        current_times,
+        model_options=model_options,
+        seed=0,
+        n_steps=10,
+    )
+
+    # With an adjacency-only boundary guard, these changes should not prevent early-stop.
+    assert calls["langevin"] == 2
+    assert calls["with_score"] == 2
+    assert calls["without_score"] == 0
+
+
 def test_semantic_stop_distance_fn_dispatch_variants() -> None:
     # 1. Standard 3-arg
     def dist3(p, c, ctx):
@@ -708,7 +910,6 @@ def dist_varkw(p, c, **kwargs):
         model_options = {
             "lanpaint_semantic_stop": {
                 "threshold": 1e-6,
-                "min_steps": 1,
                 "patience": 1,
                 "distance_fn": fn,
             }

From a895d2a7581d9dd7d7c1d1e10766c60085886e2c Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Sat, 24 Jan 2026 12:43:23 +0800
Subject: [PATCH 20/33] Tighten semantic early-stop robustness

---
 src/LanPaint/earlystop.py            |  2 +-
 tests/test_lanpaint_semantic_stop.py | 47 ++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/src/LanPaint/earlystop.py b/src/LanPaint/earlystop.py
index d526433..1d76ef1 100644
--- a/src/LanPaint/earlystop.py
+++ b/src/LanPaint/earlystop.py
@@ -85,7 +85,7 @@ def from_options(
         if enabled_early_stop:
             try:
                 abt_val = float(torch.mean(abt).item())
-            except Exception:
+            except (TypeError, ValueError):
                 abt_val = 0.0
 
             threshold_eff = threshold * _abt_scale(abt_val)
diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
index d087b0f..112b8ec 100644
--- a/tests/test_lanpaint_semantic_stop.py
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -669,6 +669,53 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
     assert calls["without_score"] == 0
 
 
+def test_semantic_stop_drift_guard_prevents_premature_stop() -> 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}
+
+    threshold = 2e-5
+    delta_sq = 1e-5
+    delta = float(delta_sq**0.5)
+
+    def fake_langevin_drift(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
+
+        step = float(calls["langevin"])
+        x0 = torch.zeros_like(x_t) + (step * delta)
+
+        # Ensure the first iteration does not look stable due to x_t noise metric.
+        return x_t + (1.0 if step <= 1 else 0.0), (None, None, x0)
+
+    engine.langevin_dynamics = fake_langevin_drift  # type: ignore[method-assign]
+
+    model_options = {
+        "lanpaint_semantic_stop": {
+            "threshold": threshold,
+            "patience": 3,
+        }
+    }
+
+    x, latent_image, noise, sigma, latent_mask, _ = _inputs()
+    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
+    engine(x, latent_image, noise, sigma, latent_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
+
+
 def test_semantic_stop_does_not_stop_while_boundary_changes() -> None:
     engine = LanPaintEngine(
         _DummyModel(),

From c7962bf78405ee0a8b8407a2928d6214d50ce9bc Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Sun, 25 Jan 2026 10:57:29 +0800
Subject: [PATCH 21/33] Refactor semantic early-stop metric

---
 src/LanPaint/earlystop.py | 33 +++++++++++++++++----------------
 src/LanPaint/lanpaint.py  |  3 +++
 2 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/src/LanPaint/earlystop.py b/src/LanPaint/earlystop.py
index 1d76ef1..a74a4e1 100644
--- a/src/LanPaint/earlystop.py
+++ b/src/LanPaint/earlystop.py
@@ -43,6 +43,12 @@ def _boundary_weight(latent_mask: torch.Tensor, inpaint_weight: torch.Tensor) ->
     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.
@@ -75,6 +81,9 @@ def from_options(
 
         enabled_early_stop = (threshold > 0.0) and (patience > 0)
         patience = max(1, patience)
+        # 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 = patience + 1
         threshold_eff = threshold
         inpaint_weight = None
@@ -250,19 +259,14 @@ def step(
                 x0_cur = args[2]
 
             if x0_prev is not None and x0_cur is not None:
-                diff_sq = (x0_cur.to(dtype=torch.float32) - x0_prev.to(dtype=torch.float32)) ** 2
-                denom = torch.sum(inpaint) + 1e-12
-                dist_inpaint = (torch.sum(diff_sq * inpaint) / denom).item()
-                dist = float(dist_inpaint)
+                dist_inpaint = _weighted_mse(x0_cur, x0_prev, inpaint)
+                dist = dist_inpaint
                 if self.ring_weight is not None:
-                    ring_denom = torch.sum(self.ring_weight) + 1e-12
-                    dist_ring = (torch.sum(diff_sq * self.ring_weight) / ring_denom).item()
+                    dist_ring = _weighted_mse(x0_cur, x0_prev, self.ring_weight)
                     dist = max(float(dist_inpaint), float(dist_ring))
             else:
-                diff_sq = (x_t_after.to(dtype=torch.float32) - x_t_before.to(dtype=torch.float32)) ** 2
-                denom = torch.sum(inpaint) + 1e-12
-                dist = (torch.sum(diff_sq * inpaint) / denom).item()
-                dist_inpaint = dist
+                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
 
@@ -272,13 +276,10 @@ def step(
                 if self.x0_anchor is None:
                     self.x0_anchor = x0_cur.detach()
                 else:
-                    diff_sq = (x0_cur.to(dtype=torch.float32) - self.x0_anchor.to(dtype=torch.float32)) ** 2
-                    denom = torch.sum(inpaint) + 1e-12
-                    drift_inpaint = (torch.sum(diff_sq * inpaint) / denom).item()
-                    dist_drift = float(drift_inpaint)
+                    drift_inpaint = _weighted_mse(x0_cur, self.x0_anchor, inpaint)
+                    dist_drift = drift_inpaint
                     if self.ring_weight is not None:
-                        ring_denom = torch.sum(self.ring_weight) + 1e-12
-                        drift_ring = (torch.sum(diff_sq * self.ring_weight) / ring_denom).item()
+                        drift_ring = _weighted_mse(x0_cur, self.x0_anchor, self.ring_weight)
                         dist_drift = max(float(drift_inpaint), float(drift_ring))
                     dist = max(float(dist), float(dist_drift))
             else:
diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index 0bc99d3..04ce960 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -192,6 +192,7 @@ def run_damped(x_t, args):
                 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
+            # args is (v, C, x0) for the next inner step.
             return x_t, (v, C, x0)
 
         def run_overdamped(x_t, args):
@@ -205,6 +206,7 @@ def run_overdamped(x_t, args):
                 x_t = x_t + (C_new - C) * dt
                 x_t = advance_time_overdamped(x_t, dt / 2, A, C, D)
                 C = C_new
+            # args is (v, C, x0); v is None in the overdamped fallback.
             return x_t, (None, C, x0)
 
         try:
@@ -220,6 +222,7 @@ def run_overdamped(x_t, args):
         except Exception:
             x_t, (v, C, x0) = run_overdamped(x_t, args)
 
+        # args is (v, C, x0); v can be None if we fell back to the overdamped update.
         return x_t, (v, C, x0)
 
     def prepare_step_size(self, current_times, step_size, sigma_x, sigma_y):

From 8f52a1d37845ac869b94e4333cc8252ab2c94557 Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Sun, 25 Jan 2026 20:32:51 +0800
Subject: [PATCH 22/33] Use NamedTuple for Langevin state

---
 src/LanPaint/earlystop.py    | 11 +++++++++--
 src/LanPaint/lanpaint.py     | 25 ++++++++++++++-----------
 src/LanPaint/types.py        | 10 ++++++++++
 tests/test_sho_regression.py | 12 ++++++++++++
 4 files changed, 45 insertions(+), 13 deletions(-)
 create mode 100644 src/LanPaint/types.py

diff --git a/src/LanPaint/earlystop.py b/src/LanPaint/earlystop.py
index a74a4e1..ef4719e 100644
--- a/src/LanPaint/earlystop.py
+++ b/src/LanPaint/earlystop.py
@@ -3,6 +3,8 @@
 
 import torch
 
+from .types import LangevinState
+
 
 def _clamp01(val: float) -> float:
     if val <= 0.0:
@@ -253,9 +255,14 @@ def step(
             custom_dist = dist is not None
 
         if dist is None:
-            if isinstance(prev_args, tuple) and len(prev_args) >= 3:
+            if isinstance(prev_args, LangevinState):
+                x0_prev = prev_args.x0
+            elif isinstance(prev_args, tuple) and len(prev_args) >= 3:
                 x0_prev = prev_args[2]
-            if isinstance(args, tuple) and len(args) >= 3:
+
+            if isinstance(args, LangevinState):
+                x0_cur = args.x0
+            elif isinstance(args, tuple) and len(args) >= 3:
                 x0_cur = args[2]
 
             if x0_prev is not None and x0_cur is not None:
diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index 04ce960..48cd795 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -2,6 +2,7 @@
 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, EarlyStopThreshold = 0.0, EarlyStopPatience = 1, EarlyStopHook = None):
@@ -126,6 +127,9 @@ def sigma_y(self, abt):
         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) and 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)
@@ -186,44 +190,43 @@ def run_damped(x_t, args):
                 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, 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
             # args is (v, C, x0) for the next inner step.
-            return x_t, (v, C, x0)
+            return x_t, LangevinState(v, C, x0)
 
         def run_overdamped(x_t, args):
             if args is None:
                 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, 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
             # args is (v, C, x0); v is None in the overdamped fallback.
-            return x_t, (None, C, x0)
+            return x_t, LangevinState(None, C, x0)
 
         try:
-            x_t_next, (v_next, C_next, x0) = 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, x0) = run_overdamped(x_t, args)
+            x_t, state = run_overdamped(x_t, args)
 
         # args is (v, C, x0); v can be None if we fell back to the overdamped update.
-        return x_t, (v, C, x0)
+        return x_t, state
 
     def prepare_step_size(self, current_times, step_size, sigma_x, sigma_y):
         # -------------------------------------------------------------------------
diff --git a/src/LanPaint/types.py b/src/LanPaint/types.py
new file mode 100644
index 0000000..a738a85
--- /dev/null
+++ b/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]
+
diff --git a/tests/test_sho_regression.py b/tests/test_sho_regression.py
index 02102ae..958ce9f 100644
--- a/tests/test_sho_regression.py
+++ b/tests/test_sho_regression.py
@@ -29,6 +29,12 @@ def score(x):
         # Execute langevin_dynamics
         # This should try run_damped -> get NaNs -> raise ValueError -> catch -> run_overdamped
         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()
@@ -53,6 +59,12 @@ def test_langevin_dynamics_fallback_on_exception() -> None:
         # Configure dynamics to raise Exception
         mock_instance.dynamics.side_effect = RuntimeError("Simulation exploded")
         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]
         assert torch.isfinite(x_out).all()
         assert v_out is None or torch.isfinite(v_out).all()

From 06da285a09eb85ec734ac9ef6d2616a7b31b01fb Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Sun, 25 Jan 2026 23:42:04 +0800
Subject: [PATCH 23/33] Fix LangevinState backward compat

---
 src/LanPaint/lanpaint.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index 48cd795..2847db6 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -128,8 +128,12 @@ def sigma_y(self, abt):
 
     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) and len(args) >= 3:
-                args = LangevinState(args[0], args[1], args[2])
+            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)

From 114d09196549c547440f3ccc2ba09e39a51dddc2 Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 26 Jan 2026 02:50:36 +0800
Subject: [PATCH 24/33] Simplify LanPaint inner-loop glue

---
 src/LanPaint/lanpaint.py | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py
index 2847db6..f1cc265 100644
--- a/src/LanPaint/lanpaint.py
+++ b/src/LanPaint/lanpaint.py
@@ -74,7 +74,7 @@ def scale_latent_inpaint(x, sigma, noise, latent_image):
 
             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 and x_t_before is not None:
+            if stopper is not None:
                 ctx = {
                     "step": i,
                     "steps_done": i + 1,
@@ -152,7 +152,7 @@ def langevin_dynamics(self, x_t, score, mask, step_size, current_times, sigma_x=
         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, x0
         def advance_time(x_t, v, dt, Gamma, A, C, D):
@@ -266,9 +266,3 @@ def prepare_step_size(self, current_times, step_size, sigma_x, sigma_y):
         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

From cd06069b5228a609fd34500a6a13e9292dfa311f Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 26 Jan 2026 10:50:39 +0800
Subject: [PATCH 25/33] Avoid early-stop when no inpaint region

---
 src/LanPaint/earlystop.py            | 15 ++++++-----
 tests/test_lanpaint_semantic_stop.py | 39 ++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/src/LanPaint/earlystop.py b/src/LanPaint/earlystop.py
index ef4719e..2db9e41 100644
--- a/src/LanPaint/earlystop.py
+++ b/src/LanPaint/earlystop.py
@@ -104,9 +104,12 @@ def from_options(
                 enabled_early_stop = False
             else:
                 inpaint_weight = (1 - latent_mask).to(dtype=torch.float32)
-                ring_weight = _boundary_weight(latent_mask, inpaint_weight)
-                if isinstance(model_options, dict):
-                    trace = model_options.get("lanpaint_semantic_trace")
+                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
@@ -270,7 +273,7 @@ def step(
                 dist = dist_inpaint
                 if self.ring_weight is not None:
                     dist_ring = _weighted_mse(x0_cur, x0_prev, self.ring_weight)
-                    dist = max(float(dist_inpaint), float(dist_ring))
+                    dist = max(dist_inpaint, dist_ring)
             else:
                 dist_inpaint = _weighted_mse(x_t_after, x_t_before, inpaint)
                 dist = dist_inpaint
@@ -287,8 +290,8 @@ def step(
                     dist_drift = drift_inpaint
                     if self.ring_weight is not None:
                         drift_ring = _weighted_mse(x0_cur, self.x0_anchor, self.ring_weight)
-                        dist_drift = max(float(drift_inpaint), float(drift_ring))
-                    dist = max(float(dist), float(dist_drift))
+                        dist_drift = max(drift_inpaint, drift_ring)
+                    dist = max(dist, dist_drift)
             else:
                 self.x0_anchor = None
 
diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
index 112b8ec..bd30847 100644
--- a/tests/test_lanpaint_semantic_stop.py
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -669,6 +669,45 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
     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
+
+
 def test_semantic_stop_drift_guard_prevents_premature_stop() -> None:
     engine = LanPaintEngine(
         _DummyModel(),

From d73f3db2654cf7fc69fc6f97004d10fef184bc70 Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 26 Jan 2026 14:02:07 +0800
Subject: [PATCH 26/33] Document distance_fn hook

---
 src/LanPaint/earlystop.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/LanPaint/earlystop.py b/src/LanPaint/earlystop.py
index 2db9e41..a8a6d0e 100644
--- a/src/LanPaint/earlystop.py
+++ b/src/LanPaint/earlystop.py
@@ -179,6 +179,16 @@ def has_custom_distance_fn(self) -> bool:
 
     @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
 
@@ -282,7 +292,7 @@ def step(
 
         # Drift guard (only for default metric with x0_cur).
         if x0_cur is not None and not custom_dist:
-            if float(dist) <= threshold_used:
+            if dist <= threshold_used:
                 if self.x0_anchor is None:
                     self.x0_anchor = x0_cur.detach()
                 else:
@@ -295,7 +305,7 @@ def step(
             else:
                 self.x0_anchor = None
 
-        if float(dist) <= threshold_used:
+        if dist <= threshold_used:
             self.patience_counter += 1
         else:
             self.patience_counter = 0
@@ -310,7 +320,7 @@ def step(
                     "outer_step": self.bench_outer_step,
                     "bench_timestep": self.bench_timestep,
                     "inner_step": i + 1,
-                    "dist": float(dist),
+                    "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),

From 5e7d4727381ad0e51dde5c3012a3c63a6f40e02e Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 26 Jan 2026 14:39:47 +0800
Subject: [PATCH 27/33] Use LangevinState in semantic stop test mocks

---
 tests/test_lanpaint_semantic_stop.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
index bd30847..a1ce625 100644
--- a/tests/test_lanpaint_semantic_stop.py
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -1,6 +1,7 @@
 import torch
 
 from src.LanPaint.lanpaint import LanPaint as LanPaintEngine
+from src.LanPaint.types import LangevinState
 
 
 class _DummySampling:
@@ -490,7 +491,7 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         x0 = torch.zeros_like(x_t) + (step * delta)
 
         # Ensure the first iteration does not look stable due to x_t noise metric.
-        return x_t + (1.0 if step <= 1 else 0.0), (None, None, x0)
+        return x_t + (1.0 if step <= 1 else 0.0), LangevinState(v=None, C=None, x0=x0)
 
     engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
 
@@ -608,7 +609,7 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         prev_x0 = None if args is None else args[2]
         x0 = x_t.detach().clone() * 0 if prev_x0 is None else prev_x0
         x0 = x0 + delta * (1 - mask)
-        return x_t, (None, None, x0)
+        return x_t, LangevinState(v=None, C=None, x0=x0)
 
     engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
 
@@ -735,7 +736,7 @@ def fake_langevin_drift(x_t, score, mask, step_size, current_times, sigma_x=1, s
         x0 = torch.zeros_like(x_t) + (step * delta)
 
         # Ensure the first iteration does not look stable due to x_t noise metric.
-        return x_t + (1.0 if step <= 1 else 0.0), (None, None, x0)
+        return x_t + (1.0 if step <= 1 else 0.0), LangevinState(v=None, C=None, x0=x0)
 
     engine.langevin_dynamics = fake_langevin_drift  # type: ignore[method-assign]
 
@@ -813,7 +814,7 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
         # Simulate a case where the interior is stable, but the boundary keeps changing.
         step = float(calls["langevin"])
         x0 = boundary_weight.to(dtype=x_t.dtype) * (step * delta)
-        return x_t, (None, None, x0)
+        return x_t, LangevinState(v=None, C=None, x0=x0)
 
     engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
 
@@ -915,7 +916,7 @@ def fake_langevin_delta(x_t, score, mask, step_size, current_times, sigma_x=1, s
 
         step = float(calls["langevin"])
         x0 = changed_mask.to(dtype=x_t.dtype) * (step * delta)
-        return x_t, (None, None, x0)
+        return x_t, LangevinState(v=None, C=None, x0=x0)
 
     engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
 

From 6ad4a18537e124a3f4553a23f7afb7aea52fb56c Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 26 Jan 2026 18:09:21 +0800
Subject: [PATCH 28/33] Refactor earlystop to reduce duplication

---
 src/LanPaint/earlystop.py | 52 ++++++++++++++-------------------------
 1 file changed, 18 insertions(+), 34 deletions(-)

diff --git a/src/LanPaint/earlystop.py b/src/LanPaint/earlystop.py
index a8a6d0e..3ea5d8f 100644
--- a/src/LanPaint/earlystop.py
+++ b/src/LanPaint/earlystop.py
@@ -67,9 +67,7 @@ def from_options(
         default_patience: int,
         default_distance_fn: Optional[Callable[..., Any]],
     ) -> Optional["LanPaintEarlyStopper"]:
-        semantic_stop = None
-        if isinstance(model_options, dict):
-            semantic_stop = model_options.get("lanpaint_semantic_stop")
+        semantic_stop = model_options.get("lanpaint_semantic_stop") if isinstance(model_options, dict) else None
 
         threshold = float(default_threshold)
         patience = int(default_patience)
@@ -82,16 +80,12 @@ def from_options(
             distance_fn = semantic_stop.get("distance_fn", distance_fn)
 
         enabled_early_stop = (threshold > 0.0) and (patience > 0)
-        patience = max(1, patience)
         # 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 = patience + 1
+        patience_eff = max(1, patience) + 1
         threshold_eff = threshold
-        inpaint_weight = None
-        ring_weight = None
-        trace = None
-        abt_val = None
+        inpaint_weight = ring_weight = trace = abt_val = None
 
         if enabled_early_stop:
             try:
@@ -115,9 +109,7 @@ def from_options(
             return None
 
         # Pre-fetch trace keys to avoid repeated dict lookups
-        bench_case_id = None
-        bench_outer_step = None
-        bench_timestep = None
+        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")
@@ -250,11 +242,7 @@ def step(
 
         dist = None
         custom_dist = False
-        dist_inpaint = None
-        dist_ring = None
-        dist_drift = None
-        x0_prev = None
-        x0_cur = None
+        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)
@@ -265,25 +253,23 @@ def step(
                     dist = float(dist.item())
                 else:
                     dist = float(dist)
-            custom_dist = dist is not None
+        custom_dist = dist is not None
 
         if dist is None:
-            if isinstance(prev_args, LangevinState):
-                x0_prev = prev_args.x0
-            elif isinstance(prev_args, tuple) and len(prev_args) >= 3:
-                x0_prev = prev_args[2]
+            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
 
-            if isinstance(args, LangevinState):
-                x0_cur = args.x0
-            elif isinstance(args, tuple) and len(args) >= 3:
-                x0_cur = args[2]
+            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 = dist_inpaint
-                if self.ring_weight is not None:
-                    dist_ring = _weighted_mse(x0_cur, x0_prev, self.ring_weight)
-                    dist = max(dist_inpaint, dist_ring)
+                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
@@ -297,10 +283,8 @@ def step(
                     self.x0_anchor = x0_cur.detach()
                 else:
                     drift_inpaint = _weighted_mse(x0_cur, self.x0_anchor, inpaint)
-                    dist_drift = drift_inpaint
-                    if self.ring_weight is not None:
-                        drift_ring = _weighted_mse(x0_cur, self.x0_anchor, self.ring_weight)
-                        dist_drift = max(drift_inpaint, drift_ring)
+                    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

From 39c90bbf62617f526d34f0f3834287035f4be4d9 Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 26 Jan 2026 21:16:05 +0800
Subject: [PATCH 29/33] Revert benchmark workflow

---
 .github/workflows/benchmark.yml | 238 --------------------------------
 1 file changed, 238 deletions(-)
 delete mode 100644 .github/workflows/benchmark.yml

diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml
deleted file mode 100644
index df21edb..0000000
--- a/.github/workflows/benchmark.yml
+++ /dev/null
@@ -1,238 +0,0 @@
-name: LanPaintBench LPIPS (ImageNet subset)
-
-on:
-  workflow_dispatch:
-    inputs:
-      base_ref:
-        description: "Base ref to benchmark against (commit SHA or branch). Leave empty to use repo default branch."
-        required: false
-        default: ""
-      subset:
-        description: "Number of ImageNet images to benchmark"
-        required: true
-        default: "1"
-      sample_steps:
-        description: "Outer diffusion steps (guided_diffusion_e2e)"
-        required: true
-        default: "20"
-      n_steps:
-        description: "LanPaint inner steps per outer step"
-        required: true
-        default: "5"
-      seed:
-        description: "Random seed"
-        required: true
-        default: "42"
-      shuffle_buffer:
-        description: "HF streaming shuffle buffer"
-        required: true
-        default: "1024"
-      semantic_threshold:
-        description: "Enable semantic early-stop if > 0 (head_on run)"
-        required: true
-        default: "0.01"
-      semantic_patience:
-        description: "Semantic early-stop patience (head_on run)"
-        required: true
-        default: "2"
-      min_seconds:
-        description: "Fail if a benchmark runtime is below this (0 to disable)"
-        required: true
-        default: "0"
-      max_case_mean_delta:
-        description: "Fail if mean per-case metric delta exceeds this (head_on - head_off)"
-        required: true
-        default: "0.01"
-      max_case_p95_delta:
-        description: "Fail if P95 per-case metric delta exceeds this (head_on - head_off)"
-        required: true
-        default: "0.02"
-      max_case_max_delta:
-        description: "Fail if max per-case metric delta exceeds this (head_on - head_off)"
-        required: true
-        default: "0.05"
-      gallery_topk:
-        description: "How many worst-case examples to render in the gallery"
-        required: true
-        default: "8"
-
-jobs:
-  benchmark:
-    name: Compare Base vs Head
-    runs-on: ubuntu-latest
-    timeout-minutes: 360
-    env:
-      HF_HOME: ~/.cache/huggingface
-      BASE_REF: ${{ inputs.base_ref || github.event.repository.default_branch }}
-      OMP_NUM_THREADS: "1"
-      MKL_NUM_THREADS: "1"
-      LANPAINTBENCH_REF: "ea5be5d"
-      CHECKPOINT_URL: "https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion_uncond.pt"
-    steps:
-      - name: Checkout Head (Current PR)
-        uses: actions/checkout@v4
-        with:
-          path: head
-          fetch-depth: 1
-
-      - name: Checkout Base (Baseline)
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ env.BASE_REF }}
-          path: base
-          fetch-depth: 1
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.12'
-          cache: 'pip'
-
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
-          pip install lpips datasets numpy pillow blobfile
-          # We don't need to install LanPaint as a package because we use sys.path in the script
-
-      - name: Cache HF + torch + checkpoints
-        uses: actions/cache@v4
-        with:
-          path: |
-            ~/.cache/huggingface
-            ~/.cache/torch
-            ~/.cache/lanpaintbench
-          key: ${{ runner.os }}-lpips-imagenet-${{ hashFiles('head/tests/benchmark_lpips.py', 'head/tests/compare_benchmarks.py', 'head/tests/render_benchmark_gallery.py') }}
-
-      - name: Clone LanPaintBench
-        run: |
-          git clone https://github.com/scraed/LanPaintBench.git bench
-          cd bench
-          git checkout "$LANPAINTBENCH_REF"
-
-      - name: Download ImageNet-256 checkpoint
-        run: |
-          CHECKPOINT_PATH="$HOME/.cache/lanpaintbench/checkpoints/256x256_diffusion_uncond.pt"
-          mkdir -p ~/.cache/lanpaintbench/checkpoints
-          if [ ! -f "$CHECKPOINT_PATH" ]; then
-            curl -L "$CHECKPOINT_URL" -o "$CHECKPOINT_PATH"
-          fi
-          ls -lh "$CHECKPOINT_PATH"
-
-      - name: Run Baseline Benchmark
-        run: |
-          CHECKPOINT_PATH="$HOME/.cache/lanpaintbench/checkpoints/256x256_diffusion_uncond.pt"
-          # Use the benchmark script from the HEAD branch to ensure test consistency
-          python head/tests/benchmark_lpips.py \
-            --mode guided_diffusion_e2e \
-            --src base/src \
-            --lanpaintbench bench \
-            --checkpoint "$CHECKPOINT_PATH" \
-            --output baseline.json \
-            --dataset benjamin-paine/imagenet-1k-256x256 \
-            --split train \
-            --subset "${{ inputs.subset }}" \
-            --seed "${{ inputs.seed }}" \
-            --shuffle-buffer "${{ inputs.shuffle_buffer }}" \
-            --no-shuffle \
-            --mask-type box \
-            --n-steps "${{ inputs.n_steps }}" \
-            --sample-steps "${{ inputs.sample_steps }}" \
-            --device cpu \
-            --min-seconds "${{ inputs.min_seconds }}" \
-            --semantic-threshold 0.0
-
-      - name: Run Head Benchmark (early-stop OFF)
-        run: |
-          CHECKPOINT_PATH="$HOME/.cache/lanpaintbench/checkpoints/256x256_diffusion_uncond.pt"
-          python head/tests/benchmark_lpips.py \
-            --mode guided_diffusion_e2e \
-            --src head/src \
-            --lanpaintbench bench \
-            --checkpoint "$CHECKPOINT_PATH" \
-            --output head_off.json \
-            --dataset benjamin-paine/imagenet-1k-256x256 \
-            --split train \
-            --subset "${{ inputs.subset }}" \
-            --seed "${{ inputs.seed }}" \
-            --shuffle-buffer "${{ inputs.shuffle_buffer }}" \
-            --no-shuffle \
-            --mask-type box \
-            --n-steps "${{ inputs.n_steps }}" \
-            --sample-steps "${{ inputs.sample_steps }}" \
-            --device cpu \
-            --min-seconds "${{ inputs.min_seconds }}" \
-            --semantic-threshold 0.0 \
-            --save-images-dir head_off_images
-
-      - name: Run Head Benchmark (early-stop ON)
-        run: |
-          CHECKPOINT_PATH="$HOME/.cache/lanpaintbench/checkpoints/256x256_diffusion_uncond.pt"
-          python head/tests/benchmark_lpips.py \
-            --mode guided_diffusion_e2e \
-            --src head/src \
-            --lanpaintbench bench \
-            --checkpoint "$CHECKPOINT_PATH" \
-            --output head_on.json \
-            --dataset benjamin-paine/imagenet-1k-256x256 \
-            --split train \
-            --subset "${{ inputs.subset }}" \
-            --seed "${{ inputs.seed }}" \
-            --shuffle-buffer "${{ inputs.shuffle_buffer }}" \
-            --no-shuffle \
-            --mask-type box \
-            --n-steps "${{ inputs.n_steps }}" \
-            --sample-steps "${{ inputs.sample_steps }}" \
-            --device cpu \
-            --min-seconds "${{ inputs.min_seconds }}" \
-            --semantic-threshold "${{ inputs.semantic_threshold }}" \
-            --semantic-patience "${{ inputs.semantic_patience }}" \
-            --save-images-dir head_on_images
-
-      - name: Compare baseline vs head_off
-        continue-on-error: true
-        run: |
-          python head/tests/compare_benchmarks.py --baseline baseline.json --current head_off.json --max-rel-diff 0.01 --max-abs-diff 0.0
-
-      - name: Compare head_off vs head_on
-        run: |
-          python head/tests/compare_benchmarks.py \
-            --baseline head_off.json \
-            --current head_on.json \
-            --max-rel-diff 0.01 \
-            --max-abs-diff 0.0 \
-            --case-metric lpips_inpaint_ctx \
-            --max-case-mean-delta "${{ inputs.max_case_mean_delta }}" \
-            --max-case-p95-delta "${{ inputs.max_case_p95_delta }}" \
-            --max-case-max-delta "${{ inputs.max_case_max_delta }}"
-
-      - name: Compare baseline vs head_on
-        continue-on-error: true
-        run: |
-          python head/tests/compare_benchmarks.py --baseline baseline.json --current head_on.json --max-rel-diff 0.01 --max-abs-diff 0.0
-
-      - name: Render worst-case gallery (head_off vs head_on)
-        if: ${{ always() }}
-        run: |
-          python head/tests/render_benchmark_gallery.py \
-            --baseline-json head_off.json \
-            --current-json head_on.json \
-            --baseline-images head_off_images \
-            --current-images head_on_images \
-            --out-dir gallery \
-            --metric lpips_inpaint_ctx \
-            --top-k "${{ inputs.gallery_topk }}"
-
-      - name: Upload benchmark JSON
-        if: ${{ always() }}
-        uses: actions/upload-artifact@v4
-        with:
-          name: lpips-benchmark-artifacts
-          path: |
-            baseline.json
-            head_off.json
-            head_on.json
-            head_off_images
-            head_on_images
-            gallery
-          if-no-files-found: warn

From 9ddd5f613e1adcac64f891fcf0f7cd4ff443130c Mon Sep 17 00:00:00 2001
From: godnight10061 <godnight10061@users.noreply.github.com>
Date: Mon, 26 Jan 2026 21:40:36 +0800
Subject: [PATCH 30/33] Map legacy min_steps into patience

---
 src/LanPaint/earlystop.py            | 11 +++++++++++
 tests/test_lanpaint_semantic_stop.py |  8 ++++----
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/LanPaint/earlystop.py b/src/LanPaint/earlystop.py
index 3ea5d8f..b1fa903 100644
--- a/src/LanPaint/earlystop.py
+++ b/src/LanPaint/earlystop.py
@@ -79,6 +79,17 @@ def from_options(
             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
diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
index a1ce625..2eeae52 100644
--- a/tests/test_lanpaint_semantic_stop.py
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -83,7 +83,7 @@ def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
     assert calls["distance"] == 3
 
 
-def test_semantic_stop_ignores_min_steps_option() -> None:
+def test_semantic_stop_maps_min_steps_into_patience_floor() -> None:
     engine = LanPaintEngine(
         _DummyModel(),
         NSteps=10,
@@ -117,9 +117,9 @@ def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
     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)
 
-    # 'min_steps' is ignored; stopping is controlled only by patience.
-    assert calls["langevin"] == 2
-    assert calls["distance"] == 2
+    # Legacy 'min_steps' is mapped into the effective patience floor, so stopping cannot happen before 5 steps.
+    assert calls["langevin"] == 5
+    assert calls["distance"] == 5
 
 
 def test_semantic_stop_patience_is_not_abt_boosted() -> None:

From 8fe1a9d3128b4c1ae03c688e28c64e936c02b159 Mon Sep 17 00:00:00 2001
From: scraed <czhengac@connect.ust.hk>
Date: Fri, 30 Jan 2026 10:45:40 +0800
Subject: [PATCH 31/33] Move parameters to the last for backward compatibility

---
 src/LanPaint/nodes.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/LanPaint/nodes.py b/src/LanPaint/nodes.py
index 8efdc19..881bc46 100644
--- a/src/LanPaint/nodes.py
+++ b/src/LanPaint/nodes.py
@@ -371,10 +371,10 @@ def INPUT_TYPES(s):
                 "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_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."}),
-                "LanPaint_InnerPatience": ("INT", {"default": 1, "min": 1, "max": 100, "tooltip": "Number of consecutive steps below threshold required to stop."}),
                 "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."}),
+                "LanPaint_InnerPatience": ("INT", {"default": 1, "min": 1, "max": 100, "tooltip": "Number of consecutive steps below threshold required to stop."}),
                      },
                 }
 
@@ -383,7 +383,7 @@ def INPUT_TYPES(s):
 
     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_InnerThreshold=0.0, LanPaint_InnerPatience=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
@@ -571,9 +571,9 @@ def INPUT_TYPES(s):
                      "LanPaint_Friction": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 50.0, "step": 0.1, "tooltip": "Friction parameter for fast Langevin. Lower values speed convergence but may be unstable."}),
                      "LanPaint_PromptMode": (["Image First", "Prompt First"], {"tooltip": "Image First: prioritizes image quality; Prompt First: prioritizes prompt adherence."}),
                      "LanPaint_EarlyStop": ("INT", {"default": 1, "min": 0, "max": 10000, "tooltip": "Steps to stop LanPaint early, preventing irregular patterns."}),
+                     "LanPaint_Info": ("STRING", {"default": "LanPaint Custom Sampler Adv. For more info, visit https://github.com/scraed/LanPaint. If you find it useful, please give a star ⭐️!", "multiline": True}),
                      "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."}),
                      "LanPaint_InnerPatience": ("INT", {"default": 1, "min": 1, "max": 100, "tooltip": "Number of consecutive steps below threshold required to stop."}),
-                     "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}),
                     }
                }
 
@@ -584,7 +584,7 @@ def INPUT_TYPES(s):
 
     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_InnerThreshold, LanPaint_InnerPatience, 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

From 2e3a566bb4fccfacceaea9c1fe453999ed9674bd Mon Sep 17 00:00:00 2001
From: scraed <czhengac@connect.ust.hk>
Date: Fri, 30 Jan 2026 11:23:51 +0800
Subject: [PATCH 32/33] Keep a minimal set of test

---
 tests/benchmark_lpips.py             | 652 -------------------
 tests/compare_benchmarks.py          | 157 -----
 tests/render_benchmark_gallery.py    | 135 ----
 tests/test_benchmark_schedule.py     |  25 -
 tests/test_lanpaint_semantic_stop.py | 902 ---------------------------
 tests/test_sho_regression.py         |  25 -
 6 files changed, 1896 deletions(-)
 delete mode 100644 tests/benchmark_lpips.py
 delete mode 100644 tests/compare_benchmarks.py
 delete mode 100644 tests/render_benchmark_gallery.py
 delete mode 100644 tests/test_benchmark_schedule.py

diff --git a/tests/benchmark_lpips.py b/tests/benchmark_lpips.py
deleted file mode 100644
index c85f9fe..0000000
--- a/tests/benchmark_lpips.py
+++ /dev/null
@@ -1,652 +0,0 @@
-"""LanPaintBench-style LPIPS benchmark for GitHub Actions.
-
-Supports:
-- A fast dummy-mode sanity benchmark.
-- A slow e2e-ish mode using LanPaintBench's `guided_diffusion` ImageNet-256 checkpoint.
-"""
-
-import argparse
-import json
-import sys
-import time
-import warnings
-from pathlib import Path
-
-import torch
-
-
-def _make_mask(mask_type: str, height: int, width: int) -> torch.Tensor:
-    """Return a (1, 3, H, W) mask where 1=known region, 0=to-inpaint."""
-    if height <= 0 or width <= 0:
-        raise ValueError("height/width must be positive")
-
-    mask = torch.ones((1, 3, height, width), dtype=torch.float32)
-
-    if mask_type == "half":
-        mask[:, :, :, width // 2 :] = 0.0
-        return mask
-
-    if mask_type == "box":
-        h0, h1 = height // 4, (3 * height) // 4
-        w0, w1 = width // 4, (3 * width) // 4
-        mask[:, :, h0:h1, w0:w1] = 0.0
-        return mask
-
-    if mask_type == "outpaint":
-        mask = torch.zeros((1, 3, height, width), dtype=torch.float32)
-        h0, h1 = height // 4, (3 * height) // 4
-        w0, w1 = width // 4, (3 * width) // 4
-        mask[:, :, h0:h1, w0:w1] = 1.0
-        return mask
-
-    if mask_type == "checkerboard":
-        grid = max(1, min(height, width) // 16)
-        mask_2d = torch.zeros((height, width), dtype=torch.float32)
-        for row in range(height):
-            row_parity = (row // grid) % 2
-            start = 0 if row_parity == 0 else 1
-            for block in range(start, (width + grid - 1) // grid, 2):
-                c0 = block * grid
-                c1 = min(width, (block + 1) * grid)
-                mask_2d[row, c0:c1] = 1.0
-        return mask_2d.unsqueeze(0).unsqueeze(0).repeat(1, 3, 1, 1)
-
-    raise ValueError(f"unknown mask type: {mask_type}")
-
-
-class _DummySampling:
-    def noise_scaling(self, sigma, noise, latent_image):  # type: ignore[no-untyped-def]
-        # Variance exploding convention (ComfyUI/k-diffusion style).
-        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]
-        # Cheap "denoiser" so the benchmark exercises LanPaint math deterministically on CPU.
-        sigma_ = sigma.reshape((sigma.shape[0],) + (1,) * (x.ndim - 1))
-        x0 = x / (1.0 + sigma_**2)
-        return x0, x0
-
-
-def _load_images(
-    *,
-    dataset: str,
-    split: str,
-    subset: int,
-    seed: int,
-    shuffle_buffer: int,
-    shuffle: bool,
-) -> list[torch.Tensor]:
-    import numpy as np
-    from datasets import load_dataset
-
-    ds = load_dataset(dataset, split=split, streaming=True)
-    if shuffle:
-        try:
-            ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer)
-        except Exception as e:
-            warnings.warn(f"Failed to shuffle streaming dataset; continuing without shuffle: {e}")
-
-    images: list[torch.Tensor] = []
-    for idx, item in enumerate(ds):
-        if idx >= subset:
-            break
-        image = item["image"].convert("RGB")
-        arr = np.asarray(image, dtype=np.uint8)
-        img = torch.from_numpy(arr).permute(2, 0, 1).float() / 127.5 - 1.0
-        images.append(img)
-    return images
-
-
-def _run_dummy_bench(  # type: ignore[no-untyped-def]
-    LanPaintClass,
-    images: list[torch.Tensor],
-    device: torch.device,
-    *,
-    n_steps: int,
-    sigma_value: float,
-    mask_type: str,
-    seed: int,
-) -> float:
-    import lpips
-
-    lpips_fn = lpips.LPIPS(net="alex").to(device)
-
-    sigma = torch.tensor([sigma_value], device=device)
-    abt = 1.0 / (1.0 + sigma**2)
-    flow_t = torch.zeros_like(sigma)
-    current_times = (sigma, abt, flow_t)
-
-    engine = LanPaintClass(
-        Model=_DummyModel(),
-        NSteps=n_steps,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.1,
-    )
-
-    scores: list[float] = []
-    for idx, img in enumerate(images):
-        img = img.unsqueeze(0).to(device)
-        latent_image = img.clone()
-
-        height, width = img.shape[-2:]
-        latent_mask = _make_mask(mask_type, height, width).to(device)
-
-        gen = torch.Generator(device=device.type).manual_seed(seed + idx)
-        noise = torch.randn(img.shape, generator=gen, device=device)
-        x = noise * sigma.reshape((1,) + (1,) * (img.ndim - 1))
-
-        output = engine(
-            x=x,
-            latent_image=latent_image,
-            noise=noise,
-            sigma=sigma,
-            latent_mask=latent_mask,
-            current_times=current_times,
-            model_options={},
-            seed=seed,
-            n_steps=n_steps,
-        ).clamp(-1.0, 1.0)
-
-        dist = lpips_fn(output, img)
-        scores.append(float(dist.item()))
-
-    return float(sum(scores) / max(1, len(scores)))
-
-
-def _tensor_to_pil_rgb(image: torch.Tensor):  # type: ignore[no-untyped-def]
-    from PIL import Image
-
-    if image.ndim == 4:
-        image = image[0]
-    if image.ndim != 3 or image.shape[0] != 3:
-        raise ValueError(f"expected (3,H,W) or (1,3,H,W), got {tuple(image.shape)}")
-    image = image.detach().cpu().clamp(-1.0, 1.0)
-    image = (image + 1.0) * 0.5 * 255.0
-    arr = image.permute(1, 2, 0).to(torch.uint8).numpy()
-    return Image.fromarray(arr, mode="RGB")
-
-
-def _tensor_to_pil_mask(mask: torch.Tensor):  # type: ignore[no-untyped-def]
-    from PIL import Image
-
-    if mask.ndim == 4:
-        mask = mask[0]
-    if mask.ndim != 3:
-        raise ValueError(f"expected (C,H,W) or (1,C,H,W), got {tuple(mask.shape)}")
-    mask_2d = mask[0].detach().cpu().clamp(0.0, 1.0)
-    arr = (mask_2d * 255.0).to(torch.uint8).numpy()
-    return Image.fromarray(arr, mode="L")
-
-
-def _ddpm_alpha_bar(*, timesteps: int, beta_start: float, beta_end: float) -> torch.Tensor:
-    betas = torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float64)
-    alphas = 1.0 - betas
-    return torch.cumprod(alphas, dim=0)
-
-
-def _sigmas_and_timesteps(
-    *,
-    inference_steps: int,
-    train_timesteps: int,
-    beta_start: float,
-    beta_end: float,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    if inference_steps <= 0:
-        raise ValueError("inference_steps must be > 0")
-    if train_timesteps <= 1:
-        raise ValueError("train_timesteps must be > 1")
-
-    alpha_bar = _ddpm_alpha_bar(timesteps=train_timesteps, beta_start=beta_start, beta_end=beta_end)
-    t = torch.linspace(train_timesteps - 1, 0, inference_steps, dtype=torch.float64).round().to(torch.long)
-    t = torch.unique_consecutive(t)
-    if t.numel() < 2:
-        raise ValueError("inference_steps too small after rounding")
-
-    abt = alpha_bar[t].clamp(min=1e-12, max=1.0)
-    sigmas = torch.sqrt((1.0 - abt) / abt).to(torch.float32)
-    sigmas = torch.cat([sigmas, sigmas.new_zeros(1)], dim=0)
-    return sigmas, t.to(torch.long)
-
-
-class _GuidedDiffusionSampling:
-    def noise_scaling(self, sigma, noise, latent_image):  # type: ignore[no-untyped-def]
-        return latent_image + noise * sigma
-
-
-class _GuidedDiffusionModel:
-    def __init__(self, *, unet, device: torch.device) -> None:  # type: ignore[no-untyped-def]
-        self.inner_model = self
-        self.model_sampling = _GuidedDiffusionSampling()
-        self._unet = unet
-        self._device = device
-        self.call_count = 0
-
-    @torch.no_grad()
-    def __call__(self, x, sigma, model_options=None, seed=None):  # type: ignore[no-untyped-def]
-        self.call_count += 1
-        if not isinstance(model_options, dict) or "bench_timestep" not in model_options:
-            raise ValueError("model_options['bench_timestep'] is required for guided_diffusion mode")
-
-        t_int = int(model_options["bench_timestep"])
-        t = torch.full((x.shape[0],), t_int, device=self._device, dtype=torch.long)
-
-        sigma_ = sigma.reshape((sigma.shape[0],) + (1,) * (x.ndim - 1))
-        abt = 1.0 / (1.0 + sigma_**2)
-
-        x_vp = x / torch.sqrt(1.0 + sigma_**2)
-        eps = self._unet(x_vp, t)
-        eps = eps[:, :3, :, :]
-
-        sqrt_abt = torch.sqrt(abt)
-        sqrt_1m_abt = torch.sqrt(1.0 - abt)
-        x0 = (x_vp - sqrt_1m_abt * eps) / sqrt_abt
-        return x0, x0
-
-
-def _load_guided_diffusion_unet(*, lanpaintbench: Path, checkpoint: Path, device: torch.device):  # type: ignore[no-untyped-def]
-    sys.path.insert(0, str(lanpaintbench.resolve()))
-    from guided_diffusion.unet import UNetModel  # type: ignore[import-not-found]
-
-    image_size = 256
-    num_channels = 256
-    num_head_channels = 64
-    num_res_blocks = 2
-    resblock_updown = True
-    learn_sigma = True
-    class_cond = False
-    use_checkpoint = False
-    attention_resolutions = "32,16,8"
-    use_scale_shift_norm = True
-    use_fp16 = False
-
-    channel_mult = (1, 1, 2, 2, 4, 4)
-    attention_ds = []
-    for res in attention_resolutions.split(","):
-        attention_ds.append(image_size // int(res))
-
-    model = UNetModel(
-        image_size=image_size,
-        in_channels=3,
-        model_channels=num_channels,
-        out_channels=(3 if not learn_sigma else 6),
-        num_res_blocks=num_res_blocks,
-        attention_resolutions=tuple(attention_ds),
-        dropout=0.0,
-        channel_mult=channel_mult,
-        num_classes=(1000 if class_cond else None),
-        use_checkpoint=use_checkpoint,
-        use_fp16=use_fp16,
-        num_heads=1,
-        num_head_channels=num_head_channels,
-        num_heads_upsample=-1,
-        use_scale_shift_norm=use_scale_shift_norm,
-        resblock_updown=resblock_updown,
-        use_new_attention_order=False,
-    ).to(device)
-
-    state = torch.load(checkpoint, map_location=device)
-    model.load_state_dict(state)
-    model.eval()
-    return model
-
-
-def _run_guided_diffusion_e2e_bench(  # type: ignore[no-untyped-def]
-    LanPaintClass,
-    images: list[torch.Tensor],
-    device: torch.device,
-    *,
-    lanpaintbench: Path,
-    checkpoint: Path,
-    seed: int,
-    mask_type: str,
-    sample_steps: int,
-    inner_steps: int,
-    beta_start: float,
-    beta_end: float,
-    train_timesteps: int,
-    semantic_stop: dict[str, float | int] | None = None,
-    outer_cutoff_steps: int = 0,
-    semantic_trace: list[dict[str, float | int | str | None]] | None = None,
-    report_earlystop: bool = False,
-    save_images_dir: Path | None = None,
-) -> tuple[float, dict[str, float], list[dict[str, float | int | str]]]:
-    import lpips
-
-    unet = _load_guided_diffusion_unet(lanpaintbench=lanpaintbench, checkpoint=checkpoint, device=device)
-    model = _GuidedDiffusionModel(unet=unet, device=device)
-
-    sigmas, timesteps = _sigmas_and_timesteps(
-        inference_steps=sample_steps,
-        train_timesteps=train_timesteps,
-        beta_start=beta_start,
-        beta_end=beta_end,
-    )
-    sigmas = sigmas.to(device=device)
-    timesteps = timesteps.to(device=device)
-
-    lpips_fn = lpips.LPIPS(net="alex").to(device)
-
-    scores: list[float] = []
-    start = time.perf_counter()
-    cases: list[dict[str, float | int | str]] = []
-
-    outer_steps = int(sigmas.shape[0] - 1)
-    n_images = max(1, len(images))
-    model_calls_total = 0
-    inner_steps_executed_total = 0
-    inner_steps_requested_total = 0
-    inner_steps_executed_by_outer_step = [0] * outer_steps
-    inner_steps_requested_by_outer_step = [0] * outer_steps
-
-    engine = LanPaintClass(
-        Model=model,
-        NSteps=inner_steps,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.1,
-    )
-
-    if save_images_dir is not None:
-        save_images_dir.mkdir(parents=True, exist_ok=True)
-
-    for idx, img in enumerate(images):
-        case_start = time.perf_counter()
-        img = img.unsqueeze(0).to(device)
-
-        height, width = img.shape[-2:]
-        latent_mask = _make_mask(mask_type, height, width).to(device)
-        inpaint_mask = 1.0 - latent_mask
-
-        latent_image = img * latent_mask
-
-        gen = torch.Generator(device=device.type).manual_seed(seed + idx)
-        noise = torch.randn(img.shape, generator=gen, device=device)
-
-        x = latent_image + noise * sigmas[0].reshape((1,) + (1,) * (img.ndim - 1))
-
-        case_model_calls = 0
-        case_inner_steps_executed = 0
-        case_inner_steps_requested = 0
-        for step_idx in range(sigmas.shape[0] - 1):
-            sigma = sigmas[step_idx : step_idx + 1]
-            sigma_next = sigmas[step_idx + 1 : step_idx + 2]
-            t = timesteps[min(step_idx, timesteps.shape[0] - 1)].item()
-
-            abt = 1.0 / (1.0 + sigma**2)
-            flow_t = torch.sqrt(1.0 - abt) / (torch.sqrt(1.0 - abt) + torch.sqrt(abt))
-            current_times = (sigma, abt, flow_t)
-
-            remaining = outer_steps - step_idx
-            inner_steps_eff = inner_steps
-            if outer_cutoff_steps > 0 and remaining <= outer_cutoff_steps:
-                inner_steps_eff = 0
-
-            model_options: dict[str, object] = {
-                "bench_timestep": int(t),
-                "bench_case_id": int(idx),
-                "bench_outer_step": int(step_idx),
-            }
-            if semantic_stop is not None:
-                model_options["lanpaint_semantic_stop"] = semantic_stop
-            if semantic_trace is not None:
-                model_options["lanpaint_semantic_trace"] = semantic_trace
-
-            model.call_count = 0
-            x0 = engine(
-                x=x,
-                latent_image=latent_image,
-                noise=noise,
-                sigma=sigma,
-                latent_mask=latent_mask,
-                current_times=current_times,
-                model_options=model_options,
-                seed=seed,
-                n_steps=inner_steps_eff,
-            )
-            calls_this_step = int(model.call_count)
-            case_model_calls += calls_this_step
-            model_calls_total += calls_this_step
-
-            executed = max(0, calls_this_step - 1)
-            case_inner_steps_executed += executed
-            inner_steps_executed_total += executed
-            inner_steps_executed_by_outer_step[step_idx] += executed
-
-            case_inner_steps_requested += int(inner_steps_eff)
-            inner_steps_requested_total += int(inner_steps_eff)
-            inner_steps_requested_by_outer_step[step_idx] += int(inner_steps_eff)
-
-            if float(sigma_next.item()) == 0.0:
-                x = x0
-                break
-
-            sigma_ = sigma.reshape((sigma.shape[0],) + (1,) * (x.ndim - 1))
-            sigma_next_ = sigma_next.reshape((sigma_next.shape[0],) + (1,) * (x.ndim - 1))
-            d = (x - x0) / sigma_
-            x = x + d * (sigma_next_ - sigma_)
-
-        x = x.clamp(-1.0, 1.0)
-        reference = img
-
-        # Preserve context by replacing known region with reference in both tensors.
-        x_ctx = x * inpaint_mask + reference * latent_mask
-        ref_ctx = reference
-
-        dist_inpaint_ctx = lpips_fn(x_ctx, ref_ctx)
-        dist_inpaint_zero = lpips_fn(x * inpaint_mask, reference * inpaint_mask)
-        dist_full = lpips_fn(x, reference)
-
-        mse_inpaint = torch.mean(((x - reference) * inpaint_mask) ** 2).item()
-        case_seconds = time.perf_counter() - case_start
-
-        scores.append(float(dist_inpaint_zero.item()))
-        cases.append(
-            {
-                "case_id": idx,
-                "lpips_inpaint_ctx": float(dist_inpaint_ctx.item()),
-                "lpips_inpaint_zero": float(dist_inpaint_zero.item()),
-                "lpips_full": float(dist_full.item()),
-                "mse_inpaint": float(mse_inpaint),
-                "model_calls": int(case_model_calls),
-                "inner_steps_executed": int(case_inner_steps_executed),
-                "inner_steps_requested": int(case_inner_steps_requested),
-                "seconds": float(case_seconds),
-            }
-        )
-
-        if save_images_dir is not None:
-            case_dir = save_images_dir / f"{idx:04d}"
-            case_dir.mkdir(parents=True, exist_ok=True)
-            _tensor_to_pil_rgb(reference).save(case_dir / "reference.png")
-            _tensor_to_pil_rgb(latent_image).save(case_dir / "masked_input.png")
-            _tensor_to_pil_rgb(x).save(case_dir / "output.png")
-            _tensor_to_pil_mask(latent_mask).save(case_dir / "known_mask.png")
-            _tensor_to_pil_mask(inpaint_mask).save(case_dir / "inpaint_mask.png")
-
-    elapsed = time.perf_counter() - start
-    stats = {"seconds_total": elapsed, "seconds_per_image": elapsed / max(1, len(images))}
-    if report_earlystop:
-        stats["earlystop"] = {
-            "model_calls_total": int(model_calls_total),
-            "model_calls_per_image": float(model_calls_total / n_images),
-            "inner_steps_executed_total": int(inner_steps_executed_total),
-            "inner_steps_executed_per_image": float(inner_steps_executed_total / n_images),
-            "inner_steps_requested_total": int(inner_steps_requested_total),
-            "inner_steps_requested_per_image": float(inner_steps_requested_total / n_images),
-            "inner_steps_executed_by_outer_step_mean": [float(v / n_images) for v in inner_steps_executed_by_outer_step],
-            "inner_steps_requested_by_outer_step_mean": [float(v / n_images) for v in inner_steps_requested_by_outer_step],
-        }
-    return float(sum(scores) / max(1, len(scores))), stats, cases
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description="LanPaint LPIPS benchmark (ImageNet-1k-256x256 subset)")
-    parser.add_argument("--src", type=Path, required=True, help="Path to LanPaint src directory (contains LanPaint/)")
-    parser.add_argument("--output", type=Path, required=True, help="Output JSON path")
-    parser.add_argument("--mode", type=str, default="dummy", choices=["dummy", "guided_diffusion_e2e"])
-
-    parser.add_argument("--dataset", type=str, default="benjamin-paine/imagenet-1k-256x256")
-    parser.add_argument("--split", type=str, default="train")
-    parser.add_argument("--subset", type=int, default=8)
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--shuffle-buffer", type=int, default=1024)
-    parser.add_argument("--no-shuffle", action="store_true", help="Disable streaming shuffle for deterministic case selection")
-
-    parser.add_argument("--mask-type", type=str, default="box", choices=["half", "box", "outpaint", "checkerboard"])
-    parser.add_argument("--n-steps", type=int, default=5, help="LanPaint inner iterations per call (dummy mode)")
-    parser.add_argument("--sigma", type=float, default=1.0, help="Variance-exploding sigma value (dummy mode)")
-    parser.add_argument("--device", type=str, default="cpu", choices=["cpu", "cuda", "auto"])
-
-    parser.add_argument("--lanpaintbench", type=Path, default=None, help="Path to LanPaintBench checkout (guided_diffusion_e2e)")
-    parser.add_argument("--checkpoint", type=Path, default=None, help="Path to 256x256_diffusion_uncond.pt (guided_diffusion_e2e)")
-    parser.add_argument("--sample-steps", type=int, default=20, help="Outer diffusion steps (guided_diffusion_e2e)")
-    parser.add_argument("--train-timesteps", type=int, default=1000, help="Training timesteps for beta schedule (guided_diffusion_e2e)")
-    parser.add_argument("--beta-start", type=float, default=0.0001, help="DDPM beta_start (guided_diffusion_e2e)")
-    parser.add_argument("--beta-end", type=float, default=0.02, help="DDPM beta_end (guided_diffusion_e2e)")
-    parser.add_argument("--min-seconds", type=float, default=0.0, help="Fail if benchmark runtime is below this (guided_diffusion_e2e)")
-    parser.add_argument(
-        "--outer-cutoff-steps",
-        type=int,
-        default=0,
-        help="Disable LanPaint inner iterations for the last N outer steps (guided_diffusion_e2e)",
-    )
-    parser.add_argument(
-        "--semantic-threshold",
-        type=float,
-        default=0.0,
-        help="Enable LanPaint semantic early-stop if > 0 (guided_diffusion_e2e)",
-    )
-    parser.add_argument("--semantic-patience", type=int, default=1, help="Semantic early-stop patience (guided_diffusion_e2e)")
-    parser.add_argument("--report-earlystop", action="store_true", help="Include early-stop stats in output JSON (guided_diffusion_e2e)")
-    parser.add_argument("--trace-inner", type=Path, default=None, help="If set, write per-inner-step semantic trace JSON (guided_diffusion_e2e)")
-    parser.add_argument("--save-images-dir", type=Path, default=None, help="If set, save per-case images for manual review")
-    args = parser.parse_args()
-
-    if args.subset <= 0:
-        raise SystemExit("--subset must be > 0")
-    if args.n_steps <= 0:
-        raise SystemExit("--n-steps must be > 0")
-    if args.shuffle_buffer <= 0:
-        raise SystemExit("--shuffle-buffer must be > 0")
-    if args.semantic_patience <= 0:
-        raise SystemExit("--semantic-patience must be > 0")
-    if args.outer_cutoff_steps < 0:
-        raise SystemExit("--outer-cutoff-steps must be >= 0")
-
-    device_str = args.device
-    if device_str == "auto":
-        device_str = "cuda" if torch.cuda.is_available() else "cpu"
-    device = torch.device(device_str)
-
-    sys.path.insert(0, str(args.src.resolve()))
-    from LanPaint.lanpaint import LanPaint as LanPaintEngine
-
-    torch.manual_seed(args.seed)
-
-    images = _load_images(
-        dataset=args.dataset,
-        split=args.split,
-        subset=args.subset,
-        seed=args.seed,
-        shuffle_buffer=args.shuffle_buffer,
-        shuffle=not args.no_shuffle,
-    )
-
-    stats: dict[str, float] = {}
-    cases = None
-    if args.mode == "dummy":
-        score = _run_dummy_bench(
-            LanPaintEngine,
-            images,
-            device,
-            n_steps=args.n_steps,
-            sigma_value=args.sigma,
-            mask_type=args.mask_type,
-            seed=args.seed,
-        )
-    else:
-        if args.lanpaintbench is None or args.checkpoint is None:
-            raise SystemExit("--lanpaintbench and --checkpoint are required for guided_diffusion_e2e mode")
-
-        semantic_stop = None
-        if float(args.semantic_threshold) > 0.0:
-            semantic_stop = {
-                "threshold": float(args.semantic_threshold),
-                "patience": int(args.semantic_patience),
-            }
-
-        semantic_trace = [] if args.trace_inner is not None else None
-        score, stats, cases = _run_guided_diffusion_e2e_bench(
-            LanPaintEngine,
-            images,
-            device,
-            lanpaintbench=args.lanpaintbench,
-            checkpoint=args.checkpoint,
-            seed=args.seed,
-            mask_type=args.mask_type,
-            sample_steps=args.sample_steps,
-            inner_steps=args.n_steps,
-            beta_start=args.beta_start,
-            beta_end=args.beta_end,
-            train_timesteps=args.train_timesteps,
-            semantic_stop=semantic_stop,
-            outer_cutoff_steps=int(args.outer_cutoff_steps),
-            semantic_trace=semantic_trace,
-            report_earlystop=bool(args.report_earlystop),
-            save_images_dir=args.save_images_dir,
-        )
-        if args.min_seconds > 0.0 and stats.get("seconds_total", 0.0) < args.min_seconds:
-            raise SystemExit(f"benchmark too fast: {stats.get('seconds_total', 0.0):.1f}s < {args.min_seconds:.1f}s")
-        if args.trace_inner is not None and semantic_trace is not None:
-            args.trace_inner.parent.mkdir(parents=True, exist_ok=True)
-            args.trace_inner.write_text(json.dumps(semantic_trace), encoding="utf-8")
-
-    args.output.parent.mkdir(parents=True, exist_ok=True)
-    with args.output.open("w", encoding="utf-8") as f:
-        json.dump(
-            {
-                "lpips": score,
-                "stats": stats,
-                "config": {
-                    "mode": args.mode,
-                    "dataset": args.dataset,
-                    "split": args.split,
-                    "subset": args.subset,
-                    "seed": args.seed,
-                    "shuffle_buffer": args.shuffle_buffer,
-                    "no_shuffle": bool(args.no_shuffle),
-                    "mask_type": args.mask_type,
-                    "n_steps": args.n_steps,
-                    "sigma": args.sigma,
-                    "sample_steps": args.sample_steps,
-                    "train_timesteps": args.train_timesteps,
-                    "beta_start": args.beta_start,
-                    "beta_end": args.beta_end,
-                    "lanpaintbench": str(args.lanpaintbench) if args.lanpaintbench is not None else None,
-                    "checkpoint": str(args.checkpoint) if args.checkpoint is not None else None,
-                    "device": device_str,
-                    "semantic_threshold": args.semantic_threshold,
-                    "semantic_patience": args.semantic_patience,
-                    "outer_cutoff_steps": args.outer_cutoff_steps,
-                    "report_earlystop": bool(args.report_earlystop),
-                    "trace_inner": str(args.trace_inner) if args.trace_inner is not None else None,
-                },
-                "cases": cases,
-            },
-            f,
-            indent=2,
-        )
-
-    print(f"LPIPS: {score:.6f}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tests/compare_benchmarks.py b/tests/compare_benchmarks.py
deleted file mode 100644
index b417747..0000000
--- a/tests/compare_benchmarks.py
+++ /dev/null
@@ -1,157 +0,0 @@
-"""Compare two LPIPS benchmark JSON outputs (baseline vs current)."""
-
-import argparse
-import json
-import math
-import os
-import warnings
-from pathlib import Path
-
-
-def _append_step_summary(markdown: str) -> None:
-    summary_path = os.environ.get("GITHUB_STEP_SUMMARY")
-    if not summary_path:
-        return
-    with Path(summary_path).open("a", encoding="utf-8", errors="ignore") as f:
-        f.write(markdown)
-
-
-def _quantile(values: list[float], q: float) -> float:
-    if not values:
-        return float("nan")
-    if q <= 0.0:
-        return min(values)
-    if q >= 1.0:
-        return max(values)
-    ordered = sorted(values)
-    # Nearest-rank quantile.
-    idx = max(0, min(len(ordered) - 1, math.ceil(q * len(ordered)) - 1))
-    return float(ordered[idx])
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--baseline", type=Path, required=True)
-    parser.add_argument("--current", type=Path, required=True)
-    parser.add_argument("--max-rel-diff", type=float, default=0.01, help="Fail if (current-baseline)/baseline exceeds this")
-    parser.add_argument("--max-abs-diff", type=float, default=0.0, help="Fail if current-baseline exceeds this")
-    parser.add_argument(
-        "--case-metric",
-        type=str,
-        default="lpips_inpaint_ctx",
-        help="Per-case metric key to compare when 'cases' are present",
-    )
-    parser.add_argument("--max-case-mean-delta", type=float, default=None)
-    parser.add_argument("--max-case-p95-delta", type=float, default=None)
-    parser.add_argument("--max-case-max-delta", type=float, default=None)
-    args = parser.parse_args()
-
-    baseline = json.loads(args.baseline.read_text(encoding="utf-8"))
-    current = json.loads(args.current.read_text(encoding="utf-8"))
-
-    b_score = float(baseline["lpips"])
-    c_score = float(current["lpips"])
-
-    b_stats = baseline.get("stats") if isinstance(baseline, dict) else None
-    c_stats = current.get("stats") if isinstance(current, dict) else None
-    b_seconds = float(b_stats.get("seconds_total", 0.0)) if isinstance(b_stats, dict) else 0.0
-    c_seconds = float(c_stats.get("seconds_total", 0.0)) if isinstance(c_stats, dict) else 0.0
-    b_spi = float(b_stats.get("seconds_per_image", 0.0)) if isinstance(b_stats, dict) else 0.0
-    c_spi = float(c_stats.get("seconds_per_image", 0.0)) if isinstance(c_stats, dict) else 0.0
-
-    diff = c_score - b_score
-    rel_diff = diff / b_score if b_score != 0.0 else float("inf")
-
-    lines = [
-        f"Baseline LPIPS: {b_score:.6f}",
-        f"Current LPIPS:  {c_score:.6f}",
-        f"Difference:     {diff:+.6f}",
-    ]
-    if b_score != 0.0:
-        lines.append(f"Relative diff:   {rel_diff:+.2%}")
-    print("\n".join(lines))
-
-    summary = "\n".join(
-        [
-            "### LPIPS benchmark",
-            "",
-            "| Metric | Baseline | Current | Diff |",
-            "|---|---:|---:|---:|",
-            f"| LPIPS | {b_score:.6f} | {c_score:.6f} | {diff:+.6f} |",
-            f"| Runtime (s) | {b_seconds:.1f} | {c_seconds:.1f} | {(c_seconds - b_seconds):+.1f} |",
-            f"| Sec/img | {b_spi:.2f} | {c_spi:.2f} | {(c_spi - b_spi):+.2f} |",
-            "",
-        ]
-    )
-    _append_step_summary(summary)
-
-    b_cases = baseline.get("cases") if isinstance(baseline, dict) else None
-    c_cases = current.get("cases") if isinstance(current, dict) else None
-    if isinstance(b_cases, list) and isinstance(c_cases, list) and b_cases and c_cases:
-        b_by_id = {int(row["case_id"]): row for row in b_cases if isinstance(row, dict) and "case_id" in row}
-        c_by_id = {int(row["case_id"]): row for row in c_cases if isinstance(row, dict) and "case_id" in row}
-        shared_ids = sorted(set(b_by_id.keys()) & set(c_by_id.keys()))
-        if shared_ids:
-            deltas: list[float] = []
-            b_vals: list[float] = []
-            c_vals: list[float] = []
-            for case_id in shared_ids:
-                b_row = b_by_id[case_id]
-                c_row = c_by_id[case_id]
-                try:
-                    b_val = float(b_row[args.case_metric])
-                    c_val = float(c_row[args.case_metric])
-                except Exception as e:
-                    warnings.warn(f"Skipping case {case_id} due to error: {e}")
-                    continue
-                b_vals.append(b_val)
-                c_vals.append(c_val)
-                deltas.append(c_val - b_val)
-
-            if deltas:
-                mean_delta = float(sum(deltas) / len(deltas))
-                p95_delta = _quantile(deltas, 0.95)
-                max_delta = max(deltas)
-                b_mean = float(sum(b_vals) / len(b_vals))
-                c_mean = float(sum(c_vals) / len(c_vals))
-
-                print("")
-                print(f"Case metric: {args.case_metric}")
-                print(f"Baseline mean: {b_mean:.6f}")
-                print(f"Current mean:  {c_mean:.6f}")
-                print(f"Mean delta:    {mean_delta:+.6f}")
-                print(f"P95 delta:     {p95_delta:+.6f}")
-                print(f"Max delta:     {max_delta:+.6f}")
-
-                case_summary = "\n".join(
-                    [
-                        "### Per-case deltas",
-                        "",
-                        "| Metric | Mean Δ | P95 Δ | Max Δ | Cases |",
-                        "|---|---:|---:|---:|---:|",
-                        f"| `{args.case_metric}` | {mean_delta:+.6f} | {p95_delta:+.6f} | {max_delta:+.6f} | {len(deltas)} |",
-                        "",
-                    ]
-                )
-                _append_step_summary(case_summary)
-
-                if args.max_case_mean_delta is not None and mean_delta > float(args.max_case_mean_delta):
-                    print("FAIL: per-case mean delta exceeds threshold")
-                    raise SystemExit(1)
-                if args.max_case_p95_delta is not None and p95_delta > float(args.max_case_p95_delta):
-                    print("FAIL: per-case p95 delta exceeds threshold")
-                    raise SystemExit(1)
-                if args.max_case_max_delta is not None and max_delta > float(args.max_case_max_delta):
-                    print("FAIL: per-case max delta exceeds threshold")
-                    raise SystemExit(1)
-
-    if diff > args.max_abs_diff:
-        if b_score == 0.0 or rel_diff > args.max_rel_diff:
-            print("FAIL: LPIPS regression exceeds configured threshold(s)")
-            raise SystemExit(1)
-
-    print("PASS: LPIPS within configured threshold(s)")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tests/render_benchmark_gallery.py b/tests/render_benchmark_gallery.py
deleted file mode 100644
index 0875690..0000000
--- a/tests/render_benchmark_gallery.py
+++ /dev/null
@@ -1,135 +0,0 @@
-import argparse
-import json
-import warnings
-from pathlib import Path
-
-import numpy as np
-
-
-def _load_cases(path: Path) -> dict[int, dict]:
-    data = json.loads(path.read_text(encoding="utf-8"))
-    cases = data.get("cases")
-    if not isinstance(cases, list):
-        return {}
-    out: dict[int, dict] = {}
-    for row in cases:
-        if not isinstance(row, dict) or "case_id" not in row:
-            continue
-        out[int(row["case_id"])] = row
-    return out
-
-
-def _open_rgb(path: Path):  # type: ignore[no-untyped-def]
-    from PIL import Image
-
-    return Image.open(path).convert("RGB")
-
-
-def _open_mask(path: Path):  # type: ignore[no-untyped-def]
-    from PIL import Image
-
-    return Image.open(path).convert("L")
-
-
-def _diff_heatmap(off_rgb, on_rgb, inpaint_mask_l):  # type: ignore[no-untyped-def]
-    from PIL import Image
-
-    off = np.asarray(off_rgb, dtype=np.int16)
-    on = np.asarray(on_rgb, dtype=np.int16)
-    mask = (np.asarray(inpaint_mask_l, dtype=np.uint8) > 127).astype(np.uint8)
-
-    diff = np.abs(off - on).astype(np.uint8)
-    diff = diff.max(axis=2)  # (H,W)
-    diff = diff * mask
-
-    # Scale for visibility; clamp to uint8.
-    diff = np.clip(diff.astype(np.int16) * 4, 0, 255).astype(np.uint8)
-    return Image.fromarray(diff, mode="L").convert("RGB")
-
-
-def _hstack(images):  # type: ignore[no-untyped-def]
-    from PIL import Image
-
-    if not images:
-        raise ValueError("no images to stack")
-    w, h = images[0].size
-    out = Image.new("RGB", (w * len(images), h))
-    for i, im in enumerate(images):
-        out.paste(im, (i * w, 0))
-    return out
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description="Render a small visual gallery comparing two benchmark runs")
-    parser.add_argument("--baseline-json", type=Path, required=True)
-    parser.add_argument("--current-json", type=Path, required=True)
-    parser.add_argument("--baseline-images", type=Path, required=True)
-    parser.add_argument("--current-images", type=Path, required=True)
-    parser.add_argument("--out-dir", type=Path, required=True)
-    parser.add_argument("--metric", type=str, default="lpips_inpaint_ctx")
-    parser.add_argument("--top-k", type=int, default=8)
-    args = parser.parse_args()
-
-    b_cases = _load_cases(args.baseline_json)
-    c_cases = _load_cases(args.current_json)
-    shared = sorted(set(b_cases.keys()) & set(c_cases.keys()))
-    if not shared:
-        raise SystemExit("no shared cases found in JSON files")
-
-    deltas: list[tuple[float, int]] = []
-    for case_id in shared:
-        try:
-            b_val = float(b_cases[case_id][args.metric])
-            c_val = float(c_cases[case_id][args.metric])
-        except Exception as e:
-            warnings.warn(f"Skipping case {case_id} due to error: {e}")
-            continue
-        deltas.append((c_val - b_val, case_id))
-
-    if not deltas:
-        raise SystemExit(f"no comparable case metric '{args.metric}' found")
-
-    deltas.sort(reverse=True, key=lambda x: x[0])
-    selected = [case_id for _, case_id in deltas[: max(1, args.top_k)]]
-
-    args.out_dir.mkdir(parents=True, exist_ok=True)
-    rendered: list[str] = []
-
-    for case_id in selected:
-        case_name = f"{case_id:04d}"
-        b_dir = args.baseline_images / case_name
-        c_dir = args.current_images / case_name
-
-        ref = _open_rgb(b_dir / "reference.png")
-        masked = _open_rgb(b_dir / "masked_input.png")
-        off = _open_rgb(b_dir / "output.png")
-        on = _open_rgb(c_dir / "output.png")
-        inpaint_mask = _open_mask(b_dir / "inpaint_mask.png")
-
-        diff = _diff_heatmap(off, on, inpaint_mask)
-        strip = _hstack([ref, masked, off, on, diff])
-
-        out_name = f"{case_name}_compare.png"
-        strip.save(args.out_dir / out_name)
-        rendered.append(out_name)
-
-    index = "\n".join(
-        [
-            "<!doctype html>",
-            "<html><head><meta charset='utf-8'><title>LanPaint benchmark gallery</title></head>",
-            "<body>",
-            "<h2>LanPaint benchmark gallery</h2>",
-            "<p>Columns: reference | masked_input | head_off | head_on | diff(head_off vs head_on)</p>",
-            *[
-                f"<div><img src='{name}' style='image-rendering:pixelated;max-width:100%'></div><hr>"
-                for name in rendered
-            ],
-            "</body></html>",
-            "",
-        ]
-    )
-    (args.out_dir / "index.html").write_text(index, encoding="utf-8")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tests/test_benchmark_schedule.py b/tests/test_benchmark_schedule.py
deleted file mode 100644
index 661f6cc..0000000
--- a/tests/test_benchmark_schedule.py
+++ /dev/null
@@ -1,25 +0,0 @@
-import torch
-
-from tests.benchmark_lpips import _make_mask, _sigmas_and_timesteps
-
-
-def test_sigmas_schedule_descends_and_ends_at_zero() -> None:
-    sigmas, timesteps = _sigmas_and_timesteps(
-        inference_steps=20,
-        train_timesteps=1000,
-        beta_start=0.0001,
-        beta_end=0.02,
-    )
-
-    assert sigmas.numel() == timesteps.numel() + 1
-    assert sigmas[-1].item() == 0.0
-    assert torch.all(sigmas[:-1] > 0.0)
-    assert torch.all(sigmas[:-1] >= sigmas[1:])
-    assert timesteps[0].item() >= timesteps[-1].item()
-
-
-def test_make_mask_shapes() -> None:
-    mask = _make_mask("box", 64, 64)
-    assert tuple(mask.shape) == (1, 3, 64, 64)
-    assert mask.dtype == torch.float32
-
diff --git a/tests/test_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py
index 2eeae52..e96acbe 100644
--- a/tests/test_lanpaint_semantic_stop.py
+++ b/tests/test_lanpaint_semantic_stop.py
@@ -1,7 +1,6 @@
 import torch
 
 from src.LanPaint.lanpaint import LanPaint as LanPaintEngine
-from src.LanPaint.types import LangevinState
 
 
 class _DummySampling:
@@ -29,400 +28,6 @@ def _inputs():  # type: ignore[no-untyped-def]
     return x, latent_image, noise, sigma, latent_mask, current_times
 
 
-def test_semantic_stop_triggers_deterministically_at_patience() -> None:
-    engine = LanPaintEngine(
-        _DummyModel(),
-        NSteps=10,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.2,
-    )
-
-    calls = {"langevin": 0, "distance": 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
-
-    def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
-        assert ctx["step"] == calls["distance"]
-        calls["distance"] += 1
-        return 0.0
-
-    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
-
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": 1e-6,
-            "patience": 2,
-            "distance_fn": distance_fn,
-        }
-    }
-
-    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
-    assert calls["distance"] == 3
-
-
-def test_semantic_stop_maps_min_steps_into_patience_floor() -> None:
-    engine = LanPaintEngine(
-        _DummyModel(),
-        NSteps=10,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.2,
-    )
-
-    calls = {"langevin": 0, "distance": 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
-        return x_t, args
-
-    def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
-        calls["distance"] += 1
-        return 0.0
-
-    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
-
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": 1e-6,
-            "patience": 1,
-            "min_steps": 5,
-            "distance_fn": distance_fn,
-        }
-    }
-
-    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)
-
-    # Legacy 'min_steps' is mapped into the effective patience floor, so stopping cannot happen before 5 steps.
-    assert calls["langevin"] == 5
-    assert calls["distance"] == 5
-
-
-def test_semantic_stop_patience_is_not_abt_boosted() -> None:
-    engine = LanPaintEngine(
-        _DummyModel(),
-        NSteps=10,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.2,
-    )
-
-    calls = {"langevin": 0, "distance": 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
-        return x_t, args
-
-    def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
-        calls["distance"] += 1
-        return 0.0
-
-    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
-
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": 1e-6,
-            "patience": 1,
-            "distance_fn": distance_fn,
-        }
-    }
-
-    x, latent_image, noise, sigma, latent_mask, _ = _inputs()
-
-    calls["langevin"] = 0
-    calls["distance"] = 0
-    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
-    engine(x, latent_image, noise, sigma, latent_mask, current_times, model_options=model_options, seed=0, n_steps=10)
-    assert calls["langevin"] == 2
-    assert calls["distance"] == 2
-
-    calls["langevin"] = 0
-    calls["distance"] = 0
-    current_times = (sigma, torch.tensor([0.99]), torch.tensor([0.0]))
-    engine(x, latent_image, noise, sigma, latent_mask, current_times, model_options=model_options, seed=0, n_steps=10)
-    assert calls["langevin"] == 2
-    assert calls["distance"] == 2
-
-
-def test_semantic_stop_is_not_disabled_below_old_min_abt_gate() -> None:
-    engine = LanPaintEngine(
-        _DummyModel(),
-        NSteps=10,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.2,
-    )
-
-    calls = {"langevin": 0, "distance": 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
-        return x_t, args
-
-    def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
-        calls["distance"] += 1
-        return 0.0
-
-    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
-
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": 1e-6,
-            "patience": 1,
-            "distance_fn": distance_fn,
-        }
-    }
-
-    x, latent_image, noise, sigma, latent_mask, _ = _inputs()
-    current_times = (sigma, torch.tensor([0.1]), torch.tensor([0.0]))
-    engine(x, latent_image, noise, sigma, latent_mask, current_times, model_options=model_options, seed=0, n_steps=10)
-
-    assert calls["langevin"] == 2
-    assert calls["distance"] == 2
-
-
-def test_semantic_stop_distance_fn_two_arg_order_is_cur_prev() -> None:
-    engine = LanPaintEngine(
-        _DummyModel(),
-        NSteps=10,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.2,
-    )
-
-    calls = {"langevin": 0, "distance": 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
-        return x_t + 1.0, args
-
-    def distance_fn(cur_x, prev_x):  # type: ignore[no-untyped-def]
-        assert float(torch.mean(cur_x).item()) > float(torch.mean(prev_x).item())
-        calls["distance"] += 1
-        return 0.0
-
-    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
-
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": 1e-6,
-            "patience": 2,
-            "distance_fn": distance_fn,
-        }
-    }
-
-    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["distance"] == 3
-
-
-def test_semantic_stop_distance_fn_accepts_third_param_not_named_ctx() -> None:
-    engine = LanPaintEngine(
-        _DummyModel(),
-        NSteps=10,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.2,
-    )
-
-    calls = {"langevin": 0, "distance": 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
-        return x_t, args
-
-    def distance_fn(prev_x, cur_x, meta):  # type: ignore[no-untyped-def]
-        assert isinstance(meta, dict)
-        assert meta["step"] == calls["distance"]
-        calls["distance"] += 1
-        return 0.0
-
-    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
-
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": 1e-6,
-            "patience": 2,
-            "distance_fn": distance_fn,
-        }
-    }
-
-    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["distance"] == 3
-
-
-def test_semantic_stop_distance_fn_accepts_kw_only_ctx() -> None:
-    engine = LanPaintEngine(
-        _DummyModel(),
-        NSteps=10,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.2,
-    )
-
-    calls = {"langevin": 0, "distance": 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
-        return x_t, args
-
-    def distance_fn(prev_x, cur_x, *, ctx):  # type: ignore[no-untyped-def]
-        assert isinstance(ctx, dict)
-        assert ctx["step"] == calls["distance"]
-        calls["distance"] += 1
-        return 0.0
-
-    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
-
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": 1e-6,
-            "patience": 2,
-            "distance_fn": distance_fn,
-        }
-    }
-
-    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["distance"] == 3
-
-
-def test_semantic_stop_distance_fn_accepts_scalar_tensor_return() -> None:
-    engine = LanPaintEngine(
-        _DummyModel(),
-        NSteps=10,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.2,
-    )
-
-    calls = {"langevin": 0, "distance": 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
-        return x_t, args
-
-    def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
-        calls["distance"] += 1
-        return torch.tensor(0.0)
-
-    engine.langevin_dynamics = fake_langevin  # type: ignore[method-assign]
-
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": 1e-6,
-            "patience": 2,
-            "distance_fn": distance_fn,
-        }
-    }
-
-    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["distance"] == 3
-
-
-def test_semantic_stop_distance_fn_rejects_multi_element_tensor_return() -> None:
-    engine = LanPaintEngine(
-        _DummyModel(),
-        NSteps=10,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.2,
-    )
-
-    def distance_fn(prev_x, cur_x, ctx):  # type: ignore[no-untyped-def]
-        return torch.tensor([0.0, 0.0])
-
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": 1e-6,
-            "patience": 2,
-            "distance_fn": distance_fn,
-        }
-    }
-
-    x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
-
-    try:
-        engine(x, latent_image, noise, sigma, latent_mask, current_times, model_options=model_options, seed=0, n_steps=10)
-    except TypeError:
-        return
-    raise AssertionError("expected TypeError for multi-element tensor distance")
-
-
-def test_semantic_stop_fallback_wrapper_does_not_swallow_internal_typeerror() -> None:
-    engine = LanPaintEngine(
-        _DummyModel(),
-        NSteps=10,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.2,
-    )
-
-    class DistanceFn:  # type: ignore[no-untyped-def]
-        __signature__ = "broken"
-
-        def __call__(self, prev_x, cur_x, ctx):
-            raise TypeError("internal")
-
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": 1e-6,
-            "patience": 1,
-            "distance_fn": DistanceFn(),
-        }
-    }
-
-    x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
-
-    try:
-        engine(x, latent_image, noise, sigma, latent_mask, current_times, model_options=model_options, seed=0, n_steps=10)
-    except TypeError:
-        return
-    raise AssertionError("expected TypeError from distance_fn to propagate")
-
-
 def test_default_semantic_stop_triggers_at_patience_without_custom_distance_fn() -> None:
     engine = LanPaintEngine(
         _DummyModel(),
@@ -460,216 +65,6 @@ def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y
     assert calls["without_score"] == 0
 
 
-def test_semantic_stop_is_more_conservative_near_tail_abt() -> 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}
-
-    x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
-
-    # Construct an x0 delta (MSE) that is below the mid-step effective threshold
-    # but above the near-tail effective threshold. This ensures early-stop is
-    # more conservative near the tail via threshold scaling.
-    delta_sq = 1e-7
-    delta = float(delta_sq**0.5)
-
-    def fake_langevin_delta(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
-
-        step = float(calls["langevin"])
-        x0 = torch.zeros_like(x_t) + (step * delta)
-
-        # Ensure the first iteration does not look stable due to x_t noise metric.
-        return x_t + (1.0 if step <= 1 else 0.0), LangevinState(v=None, C=None, x0=x0)
-
-    engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
-
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": 1e-6,
-            "patience": 1,
-        }
-    }
-
-    calls["langevin"] = 0
-    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
-    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
-
-    calls["langevin"] = 0
-    calls["with_score"] = 0
-    calls["without_score"] = 0
-    current_times = (sigma, torch.tensor([0.99]), torch.tensor([0.0]))
-    engine(
-        x,
-        latent_image,
-        noise,
-        sigma,
-        latent_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
-
-
-def test_semantic_stop_is_disabled_for_very_early_steps() -> 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.0]), torch.tensor([0.0]))
-    engine(
-        x,
-        latent_image,
-        noise,
-        sigma,
-        latent_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
-
-
-def test_default_semantic_stop_is_mask_normalized() -> 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}
-    delta = 0.01
-
-    def fake_langevin_delta(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
-        prev_x0 = None if args is None else args[2]
-        x0 = x_t.detach().clone() * 0 if prev_x0 is None else prev_x0
-        x0 = x0 + delta * (1 - mask)
-        return x_t, LangevinState(v=None, C=None, x0=x0)
-
-    engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
-
-    # Choose a threshold that would diverge if the metric was accidentally area-scaled
-    # (sum instead of mean) over the inpaint region.
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": 0.005,
-            "patience": 1,
-        }
-    }
-
-    x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
-    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
-
-    # Small hole (unknown region) -> mask mostly 1, with a small 0 patch.
-    small_hole_mask = torch.ones_like(latent_mask)
-    small_hole_mask[:, :, 3:5, 3:5] = 0.0
-
-    calls["langevin"] = 0
-    calls["with_score"] = 0
-    calls["without_score"] = 0
-    engine(
-        x,
-        latent_image,
-        noise,
-        sigma,
-        small_hole_mask,
-        current_times,
-        model_options=model_options,
-        seed=0,
-        n_steps=10,
-    )
-    assert calls["langevin"] == 2
-    assert calls["with_score"] == 2
-    assert calls["without_score"] == 0
-
-    # Larger hole (unknown region).
-    large_hole_mask = torch.ones_like(latent_mask)
-    large_hole_mask[:, :, 1:7, 1:7] = 0.0
-
-    calls["langevin"] = 0
-    calls["with_score"] = 0
-    calls["without_score"] = 0
-    engine(
-        x,
-        latent_image,
-        noise,
-        sigma,
-        large_hole_mask,
-        current_times,
-        model_options=model_options,
-        seed=0,
-        n_steps=10,
-    )
-    assert calls["langevin"] == 2
-    assert calls["with_score"] == 2
-    assert calls["without_score"] == 0
-
-
 def test_semantic_stop_is_disabled_when_no_inpaint_region() -> None:
     engine = LanPaintEngine(
         _DummyModel(),
@@ -707,300 +102,3 @@ def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y
     assert calls["langevin"] == 10
     assert calls["with_score"] == 10
     assert calls["without_score"] == 0
-
-
-def test_semantic_stop_drift_guard_prevents_premature_stop() -> 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}
-
-    threshold = 2e-5
-    delta_sq = 1e-5
-    delta = float(delta_sq**0.5)
-
-    def fake_langevin_drift(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
-
-        step = float(calls["langevin"])
-        x0 = torch.zeros_like(x_t) + (step * delta)
-
-        # Ensure the first iteration does not look stable due to x_t noise metric.
-        return x_t + (1.0 if step <= 1 else 0.0), LangevinState(v=None, C=None, x0=x0)
-
-    engine.langevin_dynamics = fake_langevin_drift  # type: ignore[method-assign]
-
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": threshold,
-            "patience": 3,
-        }
-    }
-
-    x, latent_image, noise, sigma, latent_mask, _ = _inputs()
-    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
-    engine(x, latent_image, noise, sigma, latent_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
-
-
-def test_semantic_stop_does_not_stop_while_boundary_changes() -> None:
-    engine = LanPaintEngine(
-        _DummyModel(),
-        NSteps=10,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.2,
-    )
-
-    batch = 1
-    channels = 4
-    height = 64
-    width = 64
-
-    x = torch.zeros((batch, channels, height, width))
-    latent_image = torch.zeros_like(x)
-    noise = torch.ones_like(x)
-    sigma = torch.tensor([1.0])
-
-    # latent_mask uses LanPaint convention: 1=known region, 0=to-inpaint.
-    latent_mask = torch.ones_like(x)
-    latent_mask[:, :, 12:52, 12:52] = 0.0
-    inpaint_mask = 1.0 - latent_mask
-
-    # Construct a 4-neighbor boundary: unknown pixels adjacent to known pixels.
-    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_mask = (inpaint_mask > 0.5) & neighbor_known
-    boundary_weight = boundary_mask.to(dtype=torch.float32)
-
-    boundary_area = float(torch.sum(boundary_weight).item())
-    inpaint_area = float(torch.sum(inpaint_mask).item())
-    assert boundary_area > 0.0
-    assert boundary_area < inpaint_area
-
-    threshold = 1e-6
-    ratio = inpaint_area / boundary_area
-    delta_sq = threshold * ratio * 0.9
-    delta = float(delta_sq**0.5)
-
-    calls = {"langevin": 0, "with_score": 0, "without_score": 0}
-
-    def fake_langevin_delta(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
-
-        # Simulate a case where the interior is stable, but the boundary keeps changing.
-        step = float(calls["langevin"])
-        x0 = boundary_weight.to(dtype=x_t.dtype) * (step * delta)
-        return x_t, LangevinState(v=None, C=None, x0=x0)
-
-    engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
-
-    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": threshold,
-            "patience": 1,
-        }
-    }
-
-    engine(
-        x,
-        latent_image,
-        noise,
-        sigma,
-        latent_mask,
-        current_times,
-        model_options=model_options,
-        seed=0,
-        n_steps=10,
-    )
-
-    # If ring changes are respected, we should never stop early.
-    assert calls["langevin"] == 10
-    assert calls["with_score"] == 10
-    assert calls["without_score"] == 0
-
-    # If ring changes were ignored, inpaint-average MSE would fall below the threshold
-    # and the sampler would stop early. This test ensures we keep running the model
-    # while the boundary is still changing.
-
-
-def test_semantic_stop_boundary_guard_is_adjacency_only() -> None:
-    engine = LanPaintEngine(
-        _DummyModel(),
-        NSteps=10,
-        Friction=15.0,
-        Lambda=1.0,
-        Beta=1.0,
-        StepSize=0.2,
-    )
-
-    batch = 1
-    channels = 4
-    height = 64
-    width = 64
-
-    x = torch.zeros((batch, channels, height, width))
-    latent_image = torch.zeros_like(x)
-    noise = torch.ones_like(x)
-    sigma = torch.tensor([1.0])
-
-    # latent_mask uses LanPaint convention: 1=known region, 0=to-inpaint.
-    latent_mask = torch.ones_like(x)
-    h0, h1 = 12, 52
-    w0, w1 = 12, 52
-    latent_mask[:, :, h0:h1, w0:w1] = 0.0
-    inpaint_mask = 1.0 - latent_mask
-
-    # Pixels at distance 3 from the boundary (inside the old 5px ring, but not in the 1px adjacency boundary).
-    offset = 3
-    changed_mask = torch.zeros_like(inpaint_mask)
-    changed_mask[:, :, h0 + offset : h0 + offset + 1, w0 + offset : w1 - offset] = 1.0
-    changed_mask[:, :, h1 - offset - 1 : h1 - offset, w0 + offset : w1 - offset] = 1.0
-    changed_mask[:, :, h0 + offset + 1 : h1 - offset - 1, w0 + offset : w0 + offset + 1] = 1.0
-    changed_mask[:, :, h0 + offset + 1 : h1 - offset - 1, w1 - offset - 1 : w1 - offset] = 1.0
-
-    # Reference old "ring": all pixels within 5px of the boundary inside the hole.
-    ring_width = 5
-    ring5_mask = torch.zeros_like(inpaint_mask)
-    ring5_mask[:, :, h0 : h0 + ring_width, w0:w1] = 1.0
-    ring5_mask[:, :, h1 - ring_width : h1, w0:w1] = 1.0
-    ring5_mask[:, :, h0 + ring_width : h1 - ring_width, w0 : w0 + ring_width] = 1.0
-    ring5_mask[:, :, h0 + ring_width : h1 - ring_width, w1 - ring_width : w1] = 1.0
-
-    changed_area = float(torch.sum(changed_mask).item())
-    ring5_area = float(torch.sum(ring5_mask).item())
-    inpaint_area = float(torch.sum(inpaint_mask).item())
-    assert changed_area > 0.0
-    assert ring5_area > changed_area
-    assert inpaint_area > ring5_area
-
-    threshold = 1e-6
-    # Choose a delta so that the inpaint-average MSE falls below threshold, but the old 5px ring-average stays above.
-    delta_sq = threshold * (ring5_area / changed_area) * 1.1
-    assert (delta_sq * changed_area / inpaint_area) < threshold
-    assert (delta_sq * changed_area / ring5_area) > threshold
-    delta = float(delta_sq**0.5)
-
-    calls = {"langevin": 0, "with_score": 0, "without_score": 0}
-
-    def fake_langevin_delta(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
-
-        step = float(calls["langevin"])
-        x0 = changed_mask.to(dtype=x_t.dtype) * (step * delta)
-        return x_t, LangevinState(v=None, C=None, x0=x0)
-
-    engine.langevin_dynamics = fake_langevin_delta  # type: ignore[method-assign]
-
-    current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
-    model_options = {
-        "lanpaint_semantic_stop": {
-            "threshold": threshold,
-            "patience": 1,
-        }
-    }
-
-    engine(
-        x,
-        latent_image,
-        noise,
-        sigma,
-        latent_mask,
-        current_times,
-        model_options=model_options,
-        seed=0,
-        n_steps=10,
-    )
-
-    # With an adjacency-only boundary guard, these changes should not prevent early-stop.
-    assert calls["langevin"] == 2
-    assert calls["with_score"] == 2
-    assert calls["without_score"] == 0
-
-
-def test_semantic_stop_distance_fn_dispatch_variants() -> None:
-    # 1. Standard 3-arg
-    def dist3(p, c, ctx):
-        dist3.called = True
-        return 0.0
-
-    dist3.called = False
-
-    # 2. 3-arg with different name
-    def dist3_alt(p, c, context):
-        dist3_alt.called = True
-        return 0.0
-
-    dist3_alt.called = False
-
-    # 3. 2-arg
-    def dist2(p, c):
-        dist2.called = True
-        return 0.0
-
-    dist2.called = False
-
-    # 4. Keyword-only ctx
-    def dist_kw(p, c, *, ctx):
-        dist_kw.called = True
-        return 0.0
-
-    dist_kw.called = False
-
-    # 5. Varargs
-    def dist_varargs(p, c, *args):
-        dist_varargs.called = True
-        return 0.0
-
-    dist_varargs.called = False
-
-    # 6. Varkwargs
-    def dist_varkw(p, c, **kwargs):
-        dist_varkw.called = True
-        assert "ctx" in kwargs
-        return 0.0
-
-    dist_varkw.called = False
-
-    variants = [dist3, dist3_alt, dist2, dist_kw, dist_varargs, dist_varkw]
-
-    for fn in variants:
-        engine = LanPaintEngine(_DummyModel(), NSteps=2, Friction=15.0, Lambda=1.0, Beta=1.0, StepSize=0.2)
-        model_options = {
-            "lanpaint_semantic_stop": {
-                "threshold": 1e-6,
-                "patience": 1,
-                "distance_fn": fn,
-            }
-        }
-        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=2)
-        assert fn.called, f"Function {fn.__name__} was not called"
diff --git a/tests/test_sho_regression.py b/tests/test_sho_regression.py
index 958ce9f..cf46119 100644
--- a/tests/test_sho_regression.py
+++ b/tests/test_sho_regression.py
@@ -43,28 +43,3 @@ def score(x):
         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, 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]
-        assert torch.isfinite(x_out).all()
-        assert v_out is None or torch.isfinite(v_out).all()

From 30cfd2f2e43206a931d4d902114858b53dbe2056 Mon Sep 17 00:00:00 2001
From: scraed <czhengac@connect.ust.hk>
Date: Fri, 30 Jan 2026 11:40:02 +0800
Subject: [PATCH 33/33] add contribution info

---
 src/LanPaint/earlystop.py | 4 ++++
 src/LanPaint/nodes.py     | 8 ++++----
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/LanPaint/earlystop.py b/src/LanPaint/earlystop.py
index b1fa903..c1dc244 100644
--- a/src/LanPaint/earlystop.py
+++ b/src/LanPaint/earlystop.py
@@ -1,3 +1,7 @@
+"""
+Early Stop Logic Contributed by `https://github.com/godnight10061`.
+"""
+
 import inspect
 from typing import Any, Callable, Optional
 
diff --git a/src/LanPaint/nodes.py b/src/LanPaint/nodes.py
index 881bc46..5c70646 100644
--- a/src/LanPaint/nodes.py
+++ b/src/LanPaint/nodes.py
@@ -373,8 +373,8 @@ def INPUT_TYPES(s):
                 "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."}),
-                "LanPaint_InnerPatience": ("INT", {"default": 1, "min": 1, "max": 100, "tooltip": "Number of consecutive steps below threshold required to stop."}),
+                "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)"}),
                      },
                 }
 
@@ -572,8 +572,8 @@ def INPUT_TYPES(s):
                      "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."}),
-                     "LanPaint_InnerPatience": ("INT", {"default": 1, "min": 1, "max": 100, "tooltip": "Number of consecutive steps below threshold required to stop."}),
+                     "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)"}),
                     }
                }