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 @@ class LanPaintEarlyStopper: 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 @@ class LanPaintEarlyStopper: 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 @@ class LanPaintEarlyStopper: 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 test_default_semantic_stop_is_mask_normalized() -> None: 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(),