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] 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 @@ class LanPaint():
         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 test_semantic_stop_triggers_deterministically_at_min_steps() -> None:
 
     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