diff --git a/src/LanPaint/earlystop.py b/src/LanPaint/earlystop.py new file mode 100644 index 0000000..c1dc244 --- /dev/null +++ b/src/LanPaint/earlystop.py @@ -0,0 +1,337 @@ +""" +Early Stop Logic Contributed by `https://github.com/godnight10061`. +""" + +import inspect +from typing import Any, Callable, Optional + +import torch + +from .types import LangevinState + + +def _clamp01(val: float) -> float: + if val <= 0.0: + return 0.0 + if val >= 1.0: + return 1.0 + return val + + +def _abt_scale(abt_val: float) -> float: + """ + Smooth, parameter-free scale based on outer-step noise level. + + - 0 at abt=0/1 (disable at extreme noise / extreme tail) + - 1 at abt=0.5 (mid-schedule) + """ + abt_val = _clamp01(abt_val) + return _clamp01(4.0 * abt_val * (1.0 - abt_val)) + + +def _boundary_weight(latent_mask: torch.Tensor, inpaint_weight: torch.Tensor) -> Optional[torch.Tensor]: + """ + Return a 4-neighbor boundary weight: unknown pixels adjacent to known pixels. + + This replaces the previous dilation-based "ring" (kernel/padding) and has no tunable hyperparameters. + """ + if latent_mask.dim() != 4: + return None + + known = latent_mask > 0.5 + neighbor_known = torch.zeros_like(known) + neighbor_known[:, :, 1:, :] |= known[:, :, :-1, :] + neighbor_known[:, :, :-1, :] |= known[:, :, 1:, :] + neighbor_known[:, :, :, 1:] |= known[:, :, :, :-1] + neighbor_known[:, :, :, :-1] |= known[:, :, :, 1:] + + boundary = (~known) & neighbor_known + return boundary.to(dtype=torch.float32) * inpaint_weight + + +def _weighted_mse(t1: torch.Tensor, t2: torch.Tensor, weight: torch.Tensor) -> float: + diff_sq = (t1.to(dtype=torch.float32) - t2.to(dtype=torch.float32)) ** 2 + denom = torch.sum(weight) + 1e-12 + return float((torch.sum(diff_sq * weight) / denom).item()) + + +class LanPaintEarlyStopper: + """ + Per-step early-stop logic for LanPaint inner (Langevin) iterations. + """ + + @classmethod + def from_options( + cls, + *, + model_options: Optional[dict], + latent_mask: torch.Tensor, + abt: torch.Tensor, + default_threshold: float, + default_patience: int, + default_distance_fn: Optional[Callable[..., Any]], + ) -> Optional["LanPaintEarlyStopper"]: + semantic_stop = model_options.get("lanpaint_semantic_stop") if isinstance(model_options, dict) else None + + threshold = float(default_threshold) + patience = int(default_patience) + distance_fn = default_distance_fn + # distance_fn contract: return None (use default metric) or a scalar (Python number / 0-d (1-element) torch.Tensor) + + if isinstance(semantic_stop, dict): + threshold = float(semantic_stop.get("threshold", threshold)) + patience = int(semantic_stop.get("patience", patience)) + distance_fn = semantic_stop.get("distance_fn", distance_fn) + + # Backward compatibility: map legacy 'min_steps' to a patience floor so it is not an independent knob. + if patience > 0: + min_steps = semantic_stop.get("min_steps") + if min_steps is not None: + try: + min_steps_int = int(min_steps) + except (TypeError, ValueError): + min_steps_int = 0 + if min_steps_int > 1: + patience = max(patience, min_steps_int - 1) + + enabled_early_stop = (threshold > 0.0) and (patience > 0) + # Require N+1 consecutive stable checks: + # - the first stable step sets patience_counter to 1 + # - `patience=1` therefore stops after 2 stable steps + patience_eff = max(1, patience) + 1 + threshold_eff = threshold + inpaint_weight = ring_weight = trace = abt_val = None + + if enabled_early_stop: + try: + abt_val = float(torch.mean(abt).item()) + except (TypeError, ValueError): + abt_val = 0.0 + + threshold_eff = threshold * _abt_scale(abt_val) + if threshold_eff <= 0.0: + enabled_early_stop = False + else: + inpaint_weight = (1 - latent_mask).to(dtype=torch.float32) + if float(torch.sum(inpaint_weight).item()) < 1e-6: + enabled_early_stop = False + else: + ring_weight = _boundary_weight(latent_mask, inpaint_weight) + if isinstance(model_options, dict): + trace = model_options.get("lanpaint_semantic_trace") + + if not enabled_early_stop: + return None + + # Pre-fetch trace keys to avoid repeated dict lookups + bench_case_id = bench_outer_step = bench_timestep = None + if isinstance(trace, list) and isinstance(model_options, dict): + bench_case_id = model_options.get("bench_case_id") + bench_outer_step = model_options.get("bench_outer_step") + bench_timestep = model_options.get("bench_timestep") + + return cls( + enabled=enabled_early_stop, + threshold=threshold, + threshold_eff=threshold_eff, + patience_eff=patience_eff, + inpaint_weight=inpaint_weight, + ring_weight=ring_weight, + distance_fn=distance_fn, + trace=trace, + bench_case_id=bench_case_id, + bench_outer_step=bench_outer_step, + bench_timestep=bench_timestep, + abt_val=abt_val, + ) + + def __init__( + self, + *, + enabled: bool, + threshold: float, + threshold_eff: float, + patience_eff: int, + inpaint_weight: Optional[torch.Tensor], + ring_weight: Optional[torch.Tensor], + distance_fn: Optional[Callable[..., Any]] = None, + trace: Optional[list] = None, + bench_case_id: Any = None, + bench_outer_step: Any = None, + bench_timestep: Any = None, + abt_val: Optional[float] = None, + ) -> None: + self.enabled = bool(enabled) + self.threshold = float(threshold) + self.threshold_eff = float(threshold_eff) + self.patience_eff = int(patience_eff) + + self.inpaint_weight = inpaint_weight + self.ring_weight = ring_weight + + self.trace = trace + self.bench_case_id = bench_case_id + self.bench_outer_step = bench_outer_step + self.bench_timestep = bench_timestep + self.abt_val = abt_val + + self.patience_counter = 0 + self.x0_anchor = None + + self._dist_wrapper = self._wrap_distance_fn(distance_fn) if self.enabled else None + + @property + def has_custom_distance_fn(self) -> bool: + return self._dist_wrapper is not None + + @staticmethod + def _wrap_distance_fn(distance_fn: Optional[Callable[..., Any]]): + """ + Wrap a user-provided `distance_fn` into a normalized callable: fn(prev, cur, ctx) -> dist|None. + + Supported signatures: + - 3+ positional (or *args): `distance_fn(prev, cur, ctx)` + - explicit / **kwargs ctx: `distance_fn(prev, cur, ctx=ctx)` + - default 2-arg: `distance_fn(cur, prev)` + + Return contract: None (use default metric) or a scalar (Python number / 0-d (1-element) torch.Tensor). + """ + if not callable(distance_fn): + return None + + try: + sig = inspect.signature(distance_fn) + params = list(sig.parameters.values()) + + has_ctx_param = "ctx" in sig.parameters + has_var_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params) + has_var_pos = any(p.kind == inspect.Parameter.VAR_POSITIONAL for p in params) + + pos_params = [ + p + for p in params + if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) + ] + + if len(pos_params) >= 3 or has_var_pos: + # 3-arg positional: fn(prev, cur, ctx) + return lambda p, c, ctx: distance_fn(p, c, ctx) + if has_ctx_param or has_var_kw: + # keyword ctx: fn(prev, cur, ctx=ctx) + return lambda p, c, ctx: distance_fn(p, c, ctx=ctx) + + # Default 2-arg: fn(cur, prev) + return lambda p, c, ctx: distance_fn(c, p) + except (ValueError, TypeError): + # Fallback for built-ins or complex callables. + def fallback_wrapper(p, c, ctx): + try: + return distance_fn(p, c, ctx) + except TypeError as e: + tb = e.__traceback__ + if tb is not None and tb.tb_frame.f_code is not fallback_wrapper.__code__: + raise + return distance_fn(c, p) + + return fallback_wrapper + + def step( + self, + *, + i: int, + n_steps: int, + x_t_before: torch.Tensor, + x_t_after: torch.Tensor, + x_t_prev_for_custom: Optional[torch.Tensor], + prev_args: Any, + args: Any, + ctx: dict, + ) -> bool: + if not self.enabled: + return False + + # 'inpaint_weight' is guaranteed to be set when enabled is True in the caller. + inpaint = self.inpaint_weight + if inpaint is None: + return False + + dist = None + custom_dist = False + dist_inpaint = dist_ring = dist_drift = x0_prev = x0_cur = None + + if self._dist_wrapper is not None: + dist = self._dist_wrapper(x_t_prev_for_custom, x_t_after, ctx) + if dist is not None: + if isinstance(dist, torch.Tensor): + if dist.numel() != 1: + raise TypeError("distance_fn must return None or a scalar / 0-d (1-element) tensor") + dist = float(dist.item()) + else: + dist = float(dist) + custom_dist = dist is not None + + if dist is None: + def _get_x0(arg: Any) -> Optional[torch.Tensor]: + if isinstance(arg, LangevinState): + return arg.x0 + if isinstance(arg, tuple) and len(arg) >= 3: + return arg[2] + return None + + x0_prev = _get_x0(prev_args) + x0_cur = _get_x0(args) + + if x0_prev is not None and x0_cur is not None: + dist_inpaint = _weighted_mse(x0_cur, x0_prev, inpaint) + dist_ring = _weighted_mse(x0_cur, x0_prev, self.ring_weight) if self.ring_weight is not None else None + dist = dist_inpaint if dist_ring is None else max(dist_inpaint, dist_ring) + else: + dist_inpaint = _weighted_mse(x_t_after, x_t_before, inpaint) + dist = dist_inpaint + + threshold_used = self.threshold if custom_dist else self.threshold_eff + + # Drift guard (only for default metric with x0_cur). + if x0_cur is not None and not custom_dist: + if dist <= threshold_used: + if self.x0_anchor is None: + self.x0_anchor = x0_cur.detach() + else: + drift_inpaint = _weighted_mse(x0_cur, self.x0_anchor, inpaint) + drift_ring = _weighted_mse(x0_cur, self.x0_anchor, self.ring_weight) if self.ring_weight is not None else None + dist_drift = drift_inpaint if drift_ring is None else max(drift_inpaint, drift_ring) + dist = max(dist, dist_drift) + else: + self.x0_anchor = None + + if dist <= threshold_used: + self.patience_counter += 1 + else: + self.patience_counter = 0 + self.x0_anchor = None + + should_stop = self.patience_counter >= self.patience_eff + + if isinstance(self.trace, list): + self.trace.append( + { + "case_id": self.bench_case_id, + "outer_step": self.bench_outer_step, + "bench_timestep": self.bench_timestep, + "inner_step": i + 1, + "dist": dist, + "dist_inpaint": None if dist_inpaint is None else float(dist_inpaint), + "dist_ring": None if dist_ring is None else float(dist_ring), + "dist_drift": None if dist_drift is None else float(dist_drift), + "threshold": float(threshold_used), + "threshold_eff": float(self.threshold_eff), + "patience_counter": int(self.patience_counter), + "patience_eff": int(self.patience_eff), + "abt": None if self.abt_val is None else float(self.abt_val), + "custom_dist": bool(custom_dist), + "stopped": bool(should_stop), + } + ) + + return bool(should_stop) + diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py index 10a6344..f1cc265 100644 --- a/src/LanPaint/lanpaint.py +++ b/src/LanPaint/lanpaint.py @@ -1,9 +1,11 @@ import torch from .utils import StochasticHarmonicOscillator from functools import partial +from .earlystop import LanPaintEarlyStopper +from .types import LangevinState class LanPaint(): - def __init__(self, Model, NSteps, Friction, Lambda, Beta, StepSize, IS_FLUX = False, IS_FLOW = False): + def __init__(self, Model, NSteps, Friction, Lambda, Beta, StepSize, IS_FLUX = False, IS_FLOW = False, EarlyStopThreshold = 0.0, EarlyStopPatience = 1, EarlyStopHook = None): self.n_steps = NSteps self.chara_lamb = Lambda self.IS_FLUX = IS_FLUX @@ -13,6 +15,9 @@ 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 +56,46 @@ 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 + stopper = LanPaintEarlyStopper.from_options( + model_options=model_options if isinstance(model_options, dict) else None, + latent_mask=latent_mask, + abt=abt, + default_threshold=self.early_stop_threshold, + default_patience=self.early_stop_patience, + default_distance_fn=self.early_stop_hook, + ) + for i in range(n_steps): score_func = partial( self.score_model, y = self.latent_image, mask = latent_mask, abt = self.add_none_dims(abt), sigma = self.add_none_dims(VE_Sigma), tflow = self.add_none_dims(Flow_t), model_options = model_options, seed = seed ) + + prev_args = args + x_t_prev = x_t.detach() if (stopper is not None and stopper.has_custom_distance_fn) else None + x_t_before = x_t if (stopper is not None and stopper.enabled) else None + x_t, args = self.langevin_dynamics(x_t, score_func , latent_mask, step_size , current_times, sigma_x = self.add_none_dims(self.sigma_x(abt)), sigma_y = self.add_none_dims(self.sigma_y(abt)), args = args) + + if stopper is not None: + ctx = { + "step": i, + "steps_done": i + 1, + "n_steps": n_steps, + "mask": latent_mask, + "latent_image": self.latent_image, + "current_times": current_times, + "seed": seed, + } + if stopper.step( + i=i, + n_steps=n_steps, + x_t_before=x_t_before, + x_t_after=x_t, + x_t_prev_for_custom=x_t_prev, + prev_args=prev_args, + args=args, + ctx=ctx, + ): + break + if IS_FLUX or IS_FLOW: x = x_t / ( self.add_none_dims(abt)**0.5 + (1-self.add_none_dims(abt))**0.5 ) else: @@ -85,6 +127,13 @@ 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): + if len(args) == 2: + # Backwards compat: older state was (v, C) without x0. + args = LangevinState(args[0], args[1], None) + elif len(args) >= 3: + args = LangevinState(args[0], args[1], args[2]) # prepare the step size and time parameters with torch.autocast(device_type=x_t.device.type, dtype=torch.float32): step_sizes = self.prepare_step_size(current_times, step_size, sigma_x, sigma_y) @@ -102,11 +151,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) + x0 = x_t + score(x_t) C = (abt**0.5 * x0 - x_t )/ (1-abt) + A * x_t - return C + return C, x0 def advance_time(x_t, v, dt, Gamma, A, C, D): dtype = x_t.dtype with torch.autocast(device_type=x_t.device.type, dtype=torch.float32): @@ -143,44 +191,46 @@ 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 = args.v + C = args.C x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D) - C_new = Coef_C(x_t) + C_new, x0 = Coef_C(x_t) v = v + Gamma**0.5 * ( C_new - C) *dt x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D) C = C_new - return x_t, (v, C) + # args is (v, C, x0) for the next inner step. + return x_t, LangevinState(v, C, x0) def run_overdamped(x_t, args): if args is None: - C = Coef_C(x_t) + C, x0 = Coef_C(x_t) x_t = advance_time_overdamped(x_t, dt, A, C, D) else: - _, C = args + C = args.C x_t = advance_time_overdamped(x_t, dt / 2, A, C, D) - C_new = Coef_C(x_t) + C_new, x0 = Coef_C(x_t) x_t = x_t + (C_new - C) * dt x_t = advance_time_overdamped(x_t, dt / 2, A, C, D) C = C_new - return x_t, (None, C) + # args is (v, C, x0); v is None in the overdamped fallback. + return x_t, LangevinState(None, C, x0) try: - x_t_next, (v_next, C_next) = run_damped(x_t, args) + x_t_next, state = run_damped(x_t, args) - if torch.isnan(x_t_next).any() or torch.isnan(v_next).any(): + v_next = state.v + if torch.isnan(x_t_next).any() or (v_next is not None and torch.isnan(v_next).any()): raise ValueError("NaN detected") x_t = x_t_next - v = v_next - C = C_next - except Exception: - x_t, (v, C) = run_overdamped(x_t, args) + x_t, state = run_overdamped(x_t, args) - return x_t, (v, C) + # args is (v, C, x0); v can be None if we fell back to the overdamped update. + return x_t, state def prepare_step_size(self, current_times, step_size, sigma_x, sigma_y): # ------------------------------------------------------------------------- @@ -216,9 +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 diff --git a/src/LanPaint/nodes.py b/src/LanPaint/nodes.py index 0cff2af..5c70646 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: @@ -368,6 +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. (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)"}), }, } @@ -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_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 @@ -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: @@ -561,6 +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. (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)"}), } } @@ -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_Info="", LanPaint_InnerThreshold=0.0, LanPaint_InnerPatience=1): 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/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_lanpaint_semantic_stop.py b/tests/test_lanpaint_semantic_stop.py new file mode 100644 index 0000000..e96acbe --- /dev/null +++ b/tests/test_lanpaint_semantic_stop.py @@ -0,0 +1,104 @@ +import torch + +from src.LanPaint.lanpaint import LanPaint as LanPaintEngine + + +class _DummySampling: + def noise_scaling(self, sigma, noise, latent_image): # type: ignore[no-untyped-def] + return latent_image + noise * sigma + + +class _DummyModel: + def __init__(self) -> None: + self.inner_model = self + self.model_sampling = _DummySampling() + + def __call__(self, x, sigma, model_options=None, seed=None): # type: ignore[no-untyped-def] + return x, x + + +def _inputs(): # type: ignore[no-untyped-def] + x = torch.zeros((1, 4, 8, 8)) + latent_image = torch.zeros_like(x) + noise = torch.ones_like(x) + sigma = torch.tensor([1.0]) + + latent_mask = torch.zeros_like(x) + current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0])) + return x, latent_image, noise, sigma, latent_mask, current_times + + +def test_default_semantic_stop_triggers_at_patience_without_custom_distance_fn() -> None: + engine = LanPaintEngine( + _DummyModel(), + NSteps=10, + Friction=15.0, + Lambda=1.0, + Beta=1.0, + StepSize=0.2, + ) + + calls = {"langevin": 0, "with_score": 0, "without_score": 0} + + def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None): # type: ignore[no-untyped-def] + calls["langevin"] += 1 + if score is None: + calls["without_score"] += 1 + else: + calls["with_score"] += 1 + return x_t, args + + engine.langevin_dynamics = fake_langevin # type: ignore[method-assign] + + model_options = { + "lanpaint_semantic_stop": { + "threshold": 1e-6, + "patience": 2, + } + } + + x, latent_image, noise, sigma, latent_mask, current_times = _inputs() + engine(x, latent_image, noise, sigma, latent_mask, current_times, model_options=model_options, seed=0, n_steps=10) + + assert calls["langevin"] == 3 + assert calls["with_score"] == 3 + assert calls["without_score"] == 0 + + +def test_semantic_stop_is_disabled_when_no_inpaint_region() -> None: + engine = LanPaintEngine( + _DummyModel(), + NSteps=10, + Friction=15.0, + Lambda=1.0, + Beta=1.0, + StepSize=0.2, + ) + + calls = {"langevin": 0, "with_score": 0, "without_score": 0} + + def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None): # type: ignore[no-untyped-def] + calls["langevin"] += 1 + if score is None: + calls["without_score"] += 1 + else: + calls["with_score"] += 1 + return x_t, args + + engine.langevin_dynamics = fake_langevin # type: ignore[method-assign] + + model_options = { + "lanpaint_semantic_stop": { + "threshold": 1e-6, + "patience": 1, + } + } + + x, latent_image, noise, sigma, latent_mask, _ = _inputs() + current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0])) + no_inpaint_mask = torch.ones_like(latent_mask) + engine(x, latent_image, noise, sigma, no_inpaint_mask, current_times, model_options=model_options, seed=0, n_steps=10) + + assert calls["langevin"] == 10 + assert calls["with_score"] == 10 + assert calls["without_score"] == 0 diff --git a/tests/test_sho_regression.py b/tests/test_sho_regression.py index a601532..cf46119 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,26 +28,18 @@ 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) + assert hasattr(args_out, "v") + assert hasattr(args_out, "C") + assert hasattr(args_out, "x0") + assert args_out[0] is args_out.v + assert args_out[1] is args_out.C + assert args_out[2] is args_out.x0 + v_out = args_out[0] # Verify that SHO was initialized and dynamics called MockSHO.assert_called() mock_instance.dynamics.assert_called() # Verify result is finite (indicating fallback to overdamped logic was successful) assert torch.isfinite(x_out).all(), "Output contains NaNs, fallback failed" assert v_out is None or torch.isfinite(v_out).all() -def test_langevin_dynamics_fallback_on_exception() -> None: - """Test that langevin_dynamics falls back to overdamped dynamics if damped dynamics raises Exception.""" - torch.manual_seed(0) - lp = LanPaint(Model=MagicMock(), NSteps=10, Friction=1.0, Lambda=1.0, Beta=1.0, StepSize=0.1) - x_t = torch.randn(1, 4, 8, 8) - lp.img_dim_size = 4 - mask = torch.zeros_like(x_t) - score = lambda x: torch.zeros_like(x) - step_size = torch.tensor([0.1]) - current_times = (torch.tensor([0.5]), torch.tensor([0.5]), torch.tensor([0.5])) - with patch("src.LanPaint.lanpaint.StochasticHarmonicOscillator") as MockSHO: - mock_instance = MockSHO.return_value - # Configure dynamics to raise Exception - mock_instance.dynamics.side_effect = RuntimeError("Simulation exploded") - x_out, (v_out, C_out) = lp.langevin_dynamics(x_t, score, mask, step_size, current_times, sigma_y=1.0) - assert torch.isfinite(x_out).all() \ No newline at end of file +