Repository navigation
Refine semantic inner-loop early stop - #1
Conversation
Summary of ChangesHello @godnight10061, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a significant enhancement to the LanPaint algorithm by implementing a semantic early-stopping mechanism for its inner iterations. The primary goal is to improve efficiency and stability by intelligently halting the Langevin dynamics when the image content within the unmasked region, particularly near the inpainting boundary, has converged sufficiently. This prevents unnecessary computation and potential over-optimization in later steps, leading to faster processing and potentially better visual results. The changes also include robust benchmarking tools to evaluate the effectiveness of this new feature. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Ignored Files
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request introduces a sophisticated semantic early stopping mechanism for the inner loop of LanPaint, which is a significant enhancement. The logic is complex but appears well-thought-out, with considerations for noise levels, ring artifacts, and drift. The inclusion of extensive tests and benchmarking scripts is excellent and greatly improves the quality and maintainability of the new feature. My review includes suggestions to improve code clarity and robustness by replacing magic numbers with constants, reducing code duplication, handling exceptions more gracefully in the test scripts, and improving the way function signatures are detected.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This PR introduces a sophisticated semantic early-stopping mechanism for the inner LanPaint loop, which is a great feature for performance. The implementation is quite complex, especially the logic for handling custom distance function hooks and the default distance calculation with ring-awareness and drift guards. The PR also includes an impressive set of new tests and benchmarking scripts, which significantly improves the project's testing infrastructure.
My review focuses on improving the robustness and maintainability of the new early-stopping logic in lanpaint.py. Specifically, I've pointed out a fragile TypeError check, a potential runtime error from unsafe type casting, a misplaced import, and a suggestion to refactor a very long and complex method for better readability. The changes in other files, including the extensive new tests, look solid.
# Conflicts: # src/LanPaint/lanpaint.py
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.
Reordered the Coef_C function definition to appear before advance_time for improved code organization and readability.
add138f to
34ded92
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request is a great improvement, simplifying the semantic early-stop mechanism and reducing the number of hyperparameters. The new LanPaintEarlyStopper class is well-structured and encapsulates the logic cleanly. The changes to use a derived effective threshold and a parameter-free boundary guard are solid. The addition of comprehensive unit tests and benchmark scripts is particularly commendable, ensuring the new implementation is robust and its performance can be tracked. I have one suggestion to make the code even more robust by narrowing a broad exception catch.
| if enabled_early_stop: | ||
| try: | ||
| abt_val = float(torch.mean(abt).item()) | ||
| except Exception: |
There was a problem hiding this comment.
The use of a broad except Exception: can hide underlying issues. If an unexpected error occurs when calculating abt_val (e.g., due to an invalid abt tensor), it will be silently caught, and abt_val will default to 0.0. This would cause the early stopping feature to be silently disabled or misconfigured for that step. It's better to catch more specific exceptions to avoid masking bugs.
| except Exception: | |
| except (ValueError, TypeError): |
|
Code Review Summary:
No further change requests. |
godnight10061
left a comment
There was a problem hiding this comment.
Code Review
Summary:
- The 2-knob direction (threshold + patience) reduces tuning surface area.
- ABT-derived threshold scaling and adjacency-only boundary guard keep behavior conservative without extra parameters.
No further change requests.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a significant and well-thought-out refinement of the semantic inner-loop early stopping mechanism. The changes simplify the hyperparameter surface, introduce a more robust boundary guard, and add comprehensive unit tests and benchmarking scripts. The new LanPaintEarlyStopper class is well-designed, handling various edge cases and providing flexible distance function support. The integration into the existing LanPaint logic is also solid. I've identified a couple of minor opportunities for code simplification and cleanup to enhance readability and maintainability.
|
|
||
| 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: |
There was a problem hiding this comment.
The condition if stopper is not None and x_t_before is not None: contains a redundant check. Based on the logic on line 73, x_t_before is guaranteed to be non-None if stopper is not None. Simplifying this condition will make the code more concise and readable.
| if stopper is not None and x_t_before is not None: | |
| if stopper is not None: |
|
|
||
|
|
||
| def Coef_C(x_t): | ||
| x0 = self.x0_evalutation(x_t, score, sigma, args) |
There was a problem hiding this comment.
The x0_evalutation method is only called from this line, is very simple, and its signature includes unused parameters (sigma, args). For better code clarity and maintainability, consider inlining its logic directly here. After this change, the x0_evalutation method can be removed entirely.
| x0 = self.x0_evalutation(x_t, score, sigma, args) | |
| x0 = x_t + score(x_t) |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a refined early stopping mechanism for the inner Langevin loop, simplifying hyperparameters and improving the logic. The changes are well-structured, with the core logic encapsulated in a new LanPaintEarlyStopper class and state management improved with a LangevinState named tuple. The PR is accompanied by an impressive suite of new unit tests and benchmarking tools. My review includes a couple of suggestions for improving code clarity in the new early stopping logic and reducing code duplication in the new tests.
| 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)) |
There was a problem hiding this comment.
The float() casts on dist_inpaint and dist_ring are redundant because _weighted_mse already returns a float. This also applies to other similar max calls with float() casts in this method (e.g., for dist_drift on line 290). Removing these will improve code clarity.
| dist = max(float(dist_inpaint), float(dist_ring)) | |
| dist = max(dist_inpaint, dist_ring) |
| 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 | ||
|
|
There was a problem hiding this comment.
There's a lot of repeated setup code across the test functions in this file (e.g., creating LanPaintEngine, the calls dictionary, and mocking langevin_dynamics). Consider refactoring this common setup into a helper function. If you're using a test runner like pytest, a fixture would be a great way to handle this. This would make the tests more concise, less repetitive, and easier to maintain.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request is a significant and well-executed refactoring of the semantic early stopping logic. Moving the logic into a dedicated LanPaintEarlyStopper class with comprehensive unit tests greatly improves maintainability and clarity. The introduction of the LangevinState NamedTuple is also a great improvement for code readability. The new early stopping mechanism is more sophisticated, with features like the drift guard and a parameter-free boundary metric, while simplifying the user-facing API. The changes are well-supported by new benchmarking and comparison scripts. I have one suggestion to improve the robustness of the early stopping logic when no inpainting area is present.
| ring_weight = _boundary_weight(latent_mask, inpaint_weight) | ||
| if isinstance(model_options, dict): | ||
| trace = model_options.get("lanpaint_semantic_trace") |
There was a problem hiding this comment.
The early stopping logic is based on measuring semantic distance in the inpainting area. If no inpainting area is provided (i.e., latent_mask is all 1s), inpaint_weight will be a tensor of zeros. This causes the default distance metric (_weighted_mse) to always return 0, which will likely trigger early stopping prematurely after patience steps, even if the model is still making useful changes in the Langevin dynamics loop. It would be more robust to disable early stopping altogether if there is no inpainting region.
| ring_weight = _boundary_weight(latent_mask, inpaint_weight) | |
| if isinstance(model_options, dict): | |
| trace = model_options.get("lanpaint_semantic_trace") | |
| if torch.sum(inpaint_weight) < 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") |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This is a well-executed pull request that significantly refactors and simplifies the semantic early stopping logic. The introduction of the LanPaintEarlyStopper class and the LangevinState named tuple greatly improves code organization and readability. The new parameter-free boundary guard and ABT-scaled threshold are thoughtful simplifications. Furthermore, the addition of comprehensive unit tests and benchmarking scripts is a fantastic contribution to the project's long-term maintainability. I have a few suggestions to further refine the new earlystop.py module by removing some redundant code and improving documentation.
| return self._dist_wrapper is not None | ||
|
|
||
| @staticmethod | ||
| def _wrap_distance_fn(distance_fn: Optional[Callable[..., Any]]): |
There was a problem hiding this comment.
This function's logic for wrapping distance_fn is quite flexible but also complex, especially regarding argument order. The 2-argument case expects fn(cur, prev), while the 3-argument case expects fn(prev, cur, ctx). This difference could be surprising to users. Adding a docstring to _wrap_distance_fn that explains the dispatch logic and the expected signatures for each case would greatly improve usability and maintainability.
| 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, LangevinState): | ||
| x0_cur = args.x0 | ||
| elif isinstance(args, tuple) and len(args) >= 3: | ||
| x0_cur = args[2] |
There was a problem hiding this comment.
The elif isinstance(..., tuple) checks for both prev_args and args appear to be unreachable. The langevin_dynamics function in lanpaint.py now ensures it always returns a LangevinState object, and the calling loop in LanPaint.LanPaint maintains this type. Therefore, prev_args and args passed to this method will only ever be None or a LangevinState instance, not a tuple. This defensive code can be simplified.
| 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, LangevinState): | |
| x0_cur = args.x0 | |
| elif isinstance(args, tuple) and len(args) >= 3: | |
| x0_cur = args[2] | |
| if isinstance(prev_args, LangevinState): | |
| x0_prev = prev_args.x0 | |
| if isinstance(args, LangevinState): | |
| x0_cur = args.x0 |
|
|
||
| # Drift guard (only for default metric with x0_cur). | ||
| if x0_cur is not None and not custom_dist: | ||
| if float(dist) <= threshold_used: |
There was a problem hiding this comment.
The dist variable is guaranteed to be a float at this point due to the preceding logic that ensures it's assigned a float value. The explicit cast float(dist) is redundant and can be removed for clarity. This also applies to the check on line 298 and the value for the dist key in the trace dictionary on line 313.
| if float(dist) <= threshold_used: | |
| if dist <= threshold_used: |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This is an excellent pull request that significantly refactors and improves the semantic early stopping logic. The introduction of LanPaintEarlyStopper encapsulates the logic cleanly, and the switch to a LangevinState named tuple for state management greatly improves code clarity and maintainability. The new parameter-free boundary detection and ABT-based threshold scaling are clever simplifications. I'm particularly impressed by the extensive and thoughtful test suite and benchmarking infrastructure added in this PR, which will be invaluable for ensuring correctness and preventing regressions. The changes are well-implemented and follow the description. I have one suggestion to improve the new tests to align better with the refactored state management.
| 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) |
There was a problem hiding this comment.
In this test and several others in this file (test_default_semantic_stop_is_mask_normalized, test_semantic_stop_drift_guard_prevents_premature_stop, etc.), the mocked langevin_dynamics function returns a raw tuple for the state.
While the LanPaintEarlyStopper correctly handles this for backward compatibility, it would be beneficial for test clarity and to ensure the primary code path is exercised if these mocks returned a LangevinState object instead. This would align the tests with the new data structure introduced in this PR.
Please add from src.LanPaint.types import LangevinState at the top of the file and update this return statement. This suggestion applies to the other fake langevin_dynamics implementations in this test file as well.
| 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) |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This PR is a significant and well-executed refactoring of the semantic early-stopping mechanism in LanPaint. It simplifies the hyperparameter surface by removing several knobs and introducing a more robust, parameter-free approach for boundary detection and threshold scaling. The introduction of the LangevinState named tuple and the LanPaintEarlyStopper class greatly improves code organization and readability. The addition of comprehensive unit tests and a new benchmarking suite is a major plus, ensuring the new logic is correct and allowing for performance regression testing. My review includes one minor suggestion to refactor a small piece of duplicated code for improved maintainability. Overall, this is an excellent pull request.
| 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, LangevinState): | ||
| x0_cur = args.x0 | ||
| elif isinstance(args, tuple) and len(args) >= 3: | ||
| x0_cur = args[2] |
There was a problem hiding this comment.
The logic for extracting x0_prev and x0_cur from prev_args and args is duplicated. To improve readability and maintainability, you could extract this logic into a local helper function.
| 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, LangevinState): | |
| x0_cur = args.x0 | |
| elif isinstance(args, tuple) and len(args) >= 3: | |
| x0_cur = 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 | |
| x0_prev = _get_x0(prev_args) | |
| x0_cur = _get_x0(args) |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request refines the semantic early-stopping mechanism for LanPaint's inner Langevin iterations. It introduces a new LanPaintEarlyStopper class and a LangevinState NamedTuple, simplifying the early stop logic by reducing hyperparameters and using a parameter-free boundary guard. The changes are well-integrated into the existing codebase, with appropriate updates to the LanPaint class and ComfyUI nodes. Comprehensive unit tests and benchmark scripts have been added or updated to validate the new functionality and ensure backward compatibility. The overall approach improves code clarity, maintainability, and configurability of the early stopping feature.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a refined semantic early-stopping mechanism for the LanPaint inner loop, aiming to reduce tuning complexity and improve efficiency. The changes include a new LanPaintEarlyStopper class with configurable thresholds, patience, and support for custom distance functions. Backward compatibility for legacy min_steps and Langevin state tuples is handled gracefully. Comprehensive unit tests for the early stopping logic and new benchmark utilities have also been added, ensuring the robustness and correctness of the new feature.
| 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 |
There was a problem hiding this comment.
The fallback_wrapper for _wrap_distance_fn is a clever way to handle built-ins or complex callables where inspect.signature might fail. It correctly attempts both (p, c, ctx) and (c, p) signatures and re-raises TypeError if it's not due to argument mismatch. This adds significant robustness to custom distance function integration.
| 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 |
There was a problem hiding this comment.
The 'Drift guard' logic is a critical addition. It prevents premature stopping if the model's state is continuously drifting, even if the step-to-step change is below the threshold. This ensures stability and prevents the early stopper from getting stuck in local minima or oscillating states, which is important for the overall correctness of the early stopping mechanism.
| 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]) |
There was a problem hiding this comment.
| x0 = x_t + score(x_t) | ||
| C = (abt**0.5 * x0 - x_t )/ (1-abt) + A * x_t | ||
| return C | ||
| return C, x0 |
| if torch.isnan(x_t_next).any() or (v_next is not None and torch.isnan(v_next).any()): | ||
| raise ValueError("NaN detected") |
There was a problem hiding this comment.
| 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" |
There was a problem hiding this comment.
The test_semantic_stop_distance_fn_dispatch_variants test is excellent. It thoroughly covers various function signatures for custom distance functions, including 3-arg, 2-arg, keyword-only ctx, varargs, and varkwargs. This ensures the _wrap_distance_fn utility is robust and handles diverse user-provided callables correctly.
Summary
This PR refines LanPaint semantic inner-loop early stop to reduce tuning surface area, based on feedback in scraed#74 (comment: scraed#74 (comment)).
This is opened on the fork first so we can validate CI before proposing it upstream.
What changed
Testing
Local:
CI: