mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-06-10 12:53:56 +02:00
fix: address whole-project code review (visible all/batch, engine consolidation, I/O)
Nine findings from a high-effort project-wide review, fixed and verified (571 passed, ruff/pyright clean): Correctness: - all/batch now remove Doubao/Jimeng/Samsung visible text marks: the visible step routes through the registry (new cli._remove_visible_auto) instead of a hardcoded GeminiEngine, so they no longer leave the wordmark intact. - batch always reads the original source (dropped the out_path-reuse that re-processed already-cleaned outputs on a re-run). - img2img_runner only retries the diffusion call on the deprecated-callback TypeError; any other TypeError now propagates instead of double-running. - gemini detect/remove and the reverse-alpha engines normalize channels via a new image_io.to_bgr, fixing a grayscale/BGRA crash in the FP-gate path. - _png_late_metadata advances its cursor by the clamped length, so a malformed chunk length no longer aborts the late AI-label scan. Cleanup / efficiency: - Consolidate the ~90%-identical Doubao/Jimeng/Samsung engines into a shared config-driven _text_mark_engine.TextMarkEngine base; each engine is now a thin subclass (TextMarkConfig + test shims). Behavior is byte-exact (the three engine test suites pass unchanged). Registry adapters collapse to one _text_mark(...) row each. Gemini stays a separate engine. - scan_head is memoized per (path, size, mtime), so identify() reads the file head once instead of ~8 times. - invisible_engine post-processing decodes/encodes the output once (chained in memory) instead of 2-4 times across stages. - Remove the orphaned get_model_id_for_profile (+ CONTROLNET_PROFILE); derive the --strength help from the strength constants (strength_default_help) so it cannot drift; share the --pipeline/--strength click options; simplify the retired --auto resolver. Net -835 lines. Tests added for the registry-routed visible pass, to_bgr, the polish/model/guidance wiring, and strength_default_help. CLAUDE.md updated for the new base module, the engine/registry changes, image_io.to_bgr, and the scan_head cache. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -374,6 +374,21 @@ class TestAllCommand:
|
||||
result = runner.invoke(main, ["all", "/nonexistent/file.png"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
def test_all_visible_step_uses_registry(self, runner, sample_png, tmp_path):
|
||||
"""Regression (#1): the `all` visible step must route through the registry
|
||||
(best_auto_mark), so Doubao/Jimeng/Samsung text marks are handled -- not just
|
||||
the Gemini sparkle via a hardcoded GeminiEngine."""
|
||||
mock_cls, _mock_engine = _mock_invisible_engine()
|
||||
output = tmp_path / "clean.png"
|
||||
with (
|
||||
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
|
||||
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
|
||||
patch("remove_ai_watermarks.watermark_registry.best_auto_mark", return_value=None) as mock_best,
|
||||
):
|
||||
result = runner.invoke(main, ["all", str(sample_png), "-o", str(output)])
|
||||
assert result.exit_code == 0, result.output
|
||||
mock_best.assert_called() # the registry auto-detector drove the visible pass
|
||||
|
||||
def test_all_preserves_rgba_across_invisible_step(self, runner, tmp_path):
|
||||
"""Regression: ``all`` must keep transparency even when the invisible
|
||||
step writes a 3-channel result (as the real diffusion engine does).
|
||||
|
||||
Reference in New Issue
Block a user