From d7e4fe8835a19e8dd3b61d154f3b1ac1f8374f11 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Wed, 3 Jun 2026 18:30:39 -0700 Subject: [PATCH] feat(invisible): upscale-floor for small inputs + unsharp post-filter Two quality knobs for the SDXL invisible pass: - min_resolution floor (default 1024, --min-resolution): small inputs are upscaled to a 1024px long-side floor before diffusion, since SDXL img2img distorts on a tiny latent (a 381x512 portrait wrecks at native). The output is restored to the original input size, so it is a transparent quality boost; it adds time/memory on small inputs. 0 disables. Extends the pure _target_size helper (now cap-or-floor-or-native, min skipped on a min>max misconfig), unit-tested without a model. - unsharp post-filter (humanizer.unsharp_mask, --unsharp, opt-in default 0): applied LAST, after the GFPGAN face pass (a pre-GFPGAN sharpen would be smoothed back over), to counter the soft/over-smoothed look that diffusion + restoration leave behind (an AI tell). Pairs with --humanize (grain). Both threaded through invisible/all/batch + the module-level helper. Verified end-to-end on a 381x512 portrait: upscaled to 1024, sharpened, restored to 381x512. Co-Authored-By: Claude Opus 4.8 --- CLAUDE.md | 6 +- README.md | 8 +- src/remove_ai_watermarks/cli.py | 34 ++++++++ src/remove_ai_watermarks/humanizer.py | 32 ++++++- src/remove_ai_watermarks/invisible_engine.py | 89 ++++++++++++++------ tests/test_humanizer.py | 34 +++++++- tests/test_invisible_engine.py | 30 +++++++ 7 files changed, 197 insertions(+), 36 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index ae57724..16a6207 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -10,7 +10,7 @@ You are a **principal Python engineer** maintaining a CLI tool and library for r - `uv run remove-ai-watermarks identify ` — provenance verdict (platform + watermark inventory + confidence); `--json` for machine output, `--no-visible` to skip the cv2 sparkle detector - `uv run remove-ai-watermarks metadata --check` — inspect AI metadata (C2PA, EXIF, PNG chunks) - `uv run remove-ai-watermarks metadata --remove -o ` — strip all AI metadata -- `uv run remove-ai-watermarks batch ` — process every supported image in a directory (output defaults to `_clean/`, set with `-o`). `--mode visible|invisible|metadata|all` (default `visible`); the invisible/all path reuses the same `--strength`/`--steps`/`--pipeline`/`--device`/`--max-resolution`/`--seed`/`--hf-token` knobs as `invisible`, `--inpaint/--no-inpaint` for the visible pass, `--humanize` for the Analog Humanizer, and `--restore-faces/--no-restore-faces` + `--restore-faces-weight` for the GFPGAN face-identity post-pass +- `uv run remove-ai-watermarks batch ` — process every supported image in a directory (output defaults to `_clean/`, set with `-o`). `--mode visible|invisible|metadata|all` (default `visible`); the invisible/all path reuses the same `--strength`/`--steps`/`--pipeline`/`--device`/`--max-resolution`/`--min-resolution`/`--seed`/`--hf-token` knobs as `invisible`, `--inpaint/--no-inpaint` for the visible pass, `--humanize` for the Analog Humanizer + `--unsharp` for the final sharpening post-filter, and `--restore-faces/--no-restore-faces` + `--restore-faces-weight` for the GFPGAN face-identity post-pass ## Test and lint @@ -28,7 +28,7 @@ You are a **principal Python engineer** maintaining a CLI tool and library for r - GPU/ML modules (invisible_engine, watermark_remover) are optional — guard imports with `is_available()` checks - Optional detection extras: `detect` (imwatermark — open SD/SDXL/FLUX watermark) and `trustmark` (Adobe TrustMark decoder; pulls torch + downloads weights). Both are guarded by `is_available()` and skipped by `identify` when absent. - Optional `restore` extra (gfpgan/facexlib/basicsr): the GFPGAN face-identity post-pass (`face_restore.py`, CLI `--restore-faces`, **EXPERIMENTAL, opt-in, OFF by default**). Guarded by `face_restore.is_available()`; when enabled it auto-skips with a debug log when the extra is absent or no face is detected. numpy<2-pinned and Python-3.12-pinned (see the `face_restore.py` Key-modules bullet). -- Tests for the *model-running* paths are limited to availability checks (multi-GB downloads). But the **pure helpers inside ML-adjacent modules are unit-tested without any download** and must stay that way: `_target_size` (native-vs-downscale, `test_invisible_engine.py`) and the MPS->CPU fallback control flow via mocked pipelines (`test_img2img_runner.py`, 100% cover). Don't skip these as "ML, needs a model" — only `remove_watermark`/the diffusion bodies do. +- Tests for the *model-running* paths are limited to availability checks (multi-GB downloads). But the **pure helpers inside ML-adjacent modules are unit-tested without any download** and must stay that way: `_target_size` (native-vs-downscale-cap-vs-upscale-floor, `test_invisible_engine.py`), `humanizer.unsharp_mask` (`test_humanizer.py`), and the MPS->CPU fallback control flow via mocked pipelines (`test_img2img_runner.py`, 100% cover). Don't skip these as "ML, needs a model" — only `remove_watermark`/the diffusion bodies do. ## Key modules @@ -76,7 +76,7 @@ Who embeds what, and whether it is locally detectable (so we know which gaps are ## Known limitations -- `invisible` pipeline processes at **native resolution by default** (`max_resolution=0`), matching the hosted raiw.cc backend (fal fast-sdxl, no pre-downscale). The old forced downscale-to-1024 -> upscale-back round-trip was the main quality loss (issue #10) and is gone; at strength ~0.05 SDXL img2img does not need the ~1024 downscale. `--max-resolution N` re-introduces an opt-in long-side cap purely to bound GPU/MPS memory on very large inputs (it reintroduces the lossy round-trip). For huge images that OOM at native, tile-based diffusion is still the proper long-term fix. **Concrete MPS data points (the OOM is memory-tier-dependent, NOT a hard MPS limit):** on a ~24 GB unified-memory machine (verified 2026-05-25, 1254x1254 gpt-image SDXL, fp32) native res OOMs at the *UNet* step (peak ~17 GiB), not only the VAE decode, and the auto-fallback in `img2img_runner` reloads on CPU and finishes (slow, ~13 min) -- the output is still weight-identical and defeats SynthID, so "looks hung/crashed" on Mac is usually this CPU fallback, not a pipeline error. On a **32 GB** unified-memory machine the same default SDXL pass runs entirely on MPS with **no CPU fallback** (verified 2026-05-31, 1122x1402 gpt-image, `all`/default, ~155 s end-to-end), so 32 GB clears the native-res UNet peak that 24 GB could not. Adding `enable_vae_tiling()` alone does NOT prevent the 24 GB OOM (the peak is the UNet, not the VAE). The fast Mac workarounds for memory-constrained machines are fp16 on MPS (roughly halves memory) or `--max-resolution` to cap the long side; neither is wired as the default. The `controlnet` pipeline adds the canny ControlNet weights on top of SDXL, so its peak is a bit higher than the plain `default` pass; the same MPS->CPU fallback covers an OOM. The native-vs-downscale decision lives in the pure helper `invisible_engine._target_size(w, h, max_resolution)` (returns `None` for native, a clamped target tuple otherwise) so it is unit-tested (`tests/test_invisible_engine.py::TestTargetSize`, the #10/#15 regression guard) without loading the model -- keep that logic in the helper, don't re-inline it. +- `invisible` pipeline processes at **native resolution for inputs whose long side is >= 1024px**, and **auto-upscales smaller inputs UP to a 1024px floor** (`min_resolution=1024`, the default; `--min-resolution 0` disables) before diffusion -- SDXL img2img distorts badly on a tiny latent (a 381x512 portrait wrecks at native, the #36 follow-up), and the output is restored to the original input size so the floor is a transparent quality boost (it adds time/memory on small inputs). `max_resolution=0` (default) means no downscale cap, matching the hosted raiw.cc backend (fal fast-sdxl, no pre-downscale). The old forced downscale-to-1024 -> upscale-back round-trip for LARGE images was the main quality loss (issue #10) and is gone; at strength ~0.05 SDXL img2img does not need a downscale. **Final `--unsharp` post-filter (`humanizer.unsharp_mask`, opt-in, default 0):** applied LAST (after the GFPGAN face pass, else it would be smoothed over) to counter the soft/over-smoothed look diffusion + restoration leave (an AI tell); ~0.5-0.8 safe, higher risks halos. Pairs with `--humanize` (grain adds sensor-noise texture, unsharp adds crispness). `--max-resolution N` re-introduces an opt-in long-side cap purely to bound GPU/MPS memory on very large inputs (it reintroduces the lossy round-trip). For huge images that OOM at native, tile-based diffusion is still the proper long-term fix. **Concrete MPS data points (the OOM is memory-tier-dependent, NOT a hard MPS limit):** on a ~24 GB unified-memory machine (verified 2026-05-25, 1254x1254 gpt-image SDXL, fp32) native res OOMs at the *UNet* step (peak ~17 GiB), not only the VAE decode, and the auto-fallback in `img2img_runner` reloads on CPU and finishes (slow, ~13 min) -- the output is still weight-identical and defeats SynthID, so "looks hung/crashed" on Mac is usually this CPU fallback, not a pipeline error. On a **32 GB** unified-memory machine the same default SDXL pass runs entirely on MPS with **no CPU fallback** (verified 2026-05-31, 1122x1402 gpt-image, `all`/default, ~155 s end-to-end), so 32 GB clears the native-res UNet peak that 24 GB could not. Adding `enable_vae_tiling()` alone does NOT prevent the 24 GB OOM (the peak is the UNet, not the VAE). The fast Mac workarounds for memory-constrained machines are fp16 on MPS (roughly halves memory) or `--max-resolution` to cap the long side; neither is wired as the default. The `controlnet` pipeline adds the canny ControlNet weights on top of SDXL, so its peak is a bit higher than the plain `default` pass; the same MPS->CPU fallback covers an OOM. The native-vs-cap-vs-floor decision lives in the pure helper `invisible_engine._target_size(w, h, max_resolution, min_resolution)` (returns `None` for native, a target tuple for a downscale cap OR an upscale floor; cap takes precedence, the floor is skipped on a min>max misconfig) so it is unit-tested (`tests/test_invisible_engine.py::TestTargetSize`, the #10/#15/#36 regression guard) without loading the model -- keep that logic in the helper, don't re-inline it. - **fp16 VAE black-output fix (issue #29, 2026-05-30):** on a **CUDA/XPU fp16** backend the stock SDXL VAE overflows to NaN and the *plain* img2img path decodes to an **all-black** image (reproduced on the raiw.cc result: a 1086x1448 input -> a uniformly black 4.6 KB PNG, mean 0). `watermark_remover._load_pipeline` / `_load_controlnet_pipeline` swap in the fp16-fixed SDXL VAE (`madebyollin/sdxl-vae-fp16-fix` = `_SDXL_FP16_VAE_ID`) when `_needs_fp16_vae_fix(model_id, DEFAULT_MODEL_ID, is_fp16)` is true -- only the default SDXL checkpoint on fp16. **cpu/mps run fp32** (the stock VAE is fine there, which is why the bug never reproduces on Mac). A custom non-SDXL `model_id` keeps its own VAE (the fp16-fix VAE is SDXL-architecture-specific). The decision is a pure helper, unit-tested without a download (`tests/test_platform.py::TestFp16VaeFix`); the actual black->clean recovery needs a CUDA GPU. **Confirmed on real CUDA hardware 2026-06-03:** running `all` on a 1086x1448 OpenAI gpt-image (the #29 repro size) at fp16 produced a normal (non-black) output, so the fp16-fix VAE swap resolves the all-black decode. (It was not reproducible on this MPS machine, which runs fp32, so the verification had to happen on an NVIDIA box.) - Pyright first run is slow (2-3 min) due to ML deps (torch/diffusers/transformers stubs); full-project `uv run pyright` can stall for many minutes — scope it to changed files. - A third-party PIL plugin autoload (e.g. an HEIF/AVIF plugin) can raise a non-OSError (`ModuleNotFoundError`), not `UnidentifiedImageError`, when opening a file. Code that opens user-supplied or unknown-format files should `except Exception`, not just `OSError`/`UnidentifiedImageError`. diff --git a/README.md b/README.md index 6c3af8e..2f42f36 100644 --- a/README.md +++ b/README.md @@ -113,7 +113,7 @@ image → encode to latent space (VAE) at native resolution → decode back to pixels (VAE) ``` -- Native resolution avoids shrinking the input to 1024 px first; that down-then-up round-trip was the main quality loss (issue #10). Use `--max-resolution N` only to cap GPU/MPS memory on very large inputs. +- Large inputs run at native resolution (no down-then-up round-trip, which was the main quality loss in issue #10); use `--max-resolution N` only to cap GPU/MPS memory on very large inputs. Small inputs (long side under 1024 px) are auto-upscaled to a 1024 px floor before diffusion, because SDXL distorts on a tiny latent, and the result is restored to the original size (a transparent quality boost). Disable the floor with `--min-resolution 0`. > **Default strength is vendor-adaptive (no flag needed).** The tool reads the C2PA issuer to detect which vendor's SynthID is present and picks the strength that clears it with the least quality loss: **OpenAI gpt-image → `0.10`**, **Google Gemini → `0.15`**, **unknown source → `0.15`**. An oracle-verified June 2026 study (clean pipeline, per-image openai.com/verify or Gemini app) found OpenAI's watermark clears at `0.05` across `1024`-`1600` px (resolution-independent) while Google's is ~3x more robust and needs `0.15`. The dominant factor is the vendor, not resolution. There is no local SynthID detector, so if the oracle still reads SynthID, raise `--strength`; if you care more about preserving fine text, lower it. (Caveat: Google's `0.15` was validated on the capped `--max-resolution 1536` path; a very large native Gemini image may need more.) > @@ -277,8 +277,10 @@ remove-ai-watermarks visible image.png -o clean.png remove-ai-watermarks erase image.png --region 1640,1930,400,100 -o clean.png # Invisible watermark only (SynthID etc.) — requires GPU -remove-ai-watermarks invisible image.png -o clean.png --humanize 4.0 -# Runs at native resolution by default. On a very large image that OOMs the +remove-ai-watermarks invisible image.png -o clean.png --humanize 4.0 --unsharp 0.5 +# --humanize adds film grain, --unsharp counters the soft "AI" look (both opt-in). +# Large images run at native resolution; small ones are upscaled to a 1024 floor +# first (disable with --min-resolution 0). On a very large image that OOMs the # GPU/MPS, cap the long side: --max-resolution 2048 # Strength is vendor-adaptive by default (OpenAI 0.10 / Google 0.15); override # with --strength. To preserve text/face structure, use --pipeline controlnet diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index 7166c48..b8d0f14 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -147,6 +147,18 @@ _controlnet_scale_option = click.option( "only (EXPERIMENTAL).", ) +_min_resolution_option = click.option( + "--min-resolution", + type=int, + default=1024, + help="Upscale long side UP to this (px) before diffusion when the input is smaller, so SDXL runs " + "near 1024 (small inputs distort at native); output is restored to the input size. 0 = off. Default 1024.", +) + +_unsharp_option = click.option( + "--unsharp", type=float, default=0.0, help="Unsharp-mask sharpening strength (0 = off, typical: 0.3-0.8)." +) + def _restore_faces_options(f: Any) -> Any: """Attach the shared GFPGAN face-restoration flags to an invisible-pipeline command.""" @@ -493,6 +505,8 @@ def cmd_erase( ) @_controlnet_scale_option @_restore_faces_options +@_min_resolution_option +@_unsharp_option @click.pass_context def cmd_invisible( ctx: click.Context, @@ -505,7 +519,9 @@ def cmd_invisible( seed: int | None, hf_token: str | None, humanize: float, + unsharp: float, max_resolution: int, + min_resolution: int, controlnet_scale: float, restore_faces: bool, restore_faces_weight: float, @@ -558,7 +574,9 @@ def cmd_invisible( guidance_scale=None, seed=seed, humanize=humanize, + unsharp=unsharp, max_resolution=max_resolution, + min_resolution=min_resolution, vendor=vendor, restore_faces=restore_faces, restore_faces_weight=restore_faces_weight, @@ -738,6 +756,8 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo ) @_controlnet_scale_option @_restore_faces_options +@_min_resolution_option +@_unsharp_option @click.pass_context def cmd_all( ctx: click.Context, @@ -753,7 +773,9 @@ def cmd_all( seed: int | None, hf_token: str | None, humanize: float, + unsharp: float, max_resolution: int, + min_resolution: int, controlnet_scale: float, restore_faces: bool, restore_faces_weight: float, @@ -852,7 +874,9 @@ def cmd_all( num_inference_steps=steps, seed=seed, humanize=humanize, + unsharp=unsharp, max_resolution=max_resolution, + min_resolution=min_resolution, vendor=vendor, restore_faces=restore_faces, restore_faces_weight=restore_faces_weight, @@ -907,7 +931,9 @@ def _process_batch_image( seed: int | None, hf_token: str | None, humanize: float, + unsharp: float = 0.0, max_resolution: int = 0, + min_resolution: int = 1024, restore_faces: bool = False, restore_faces_weight: float = 0.5, ) -> None: @@ -969,7 +995,9 @@ def _process_batch_image( num_inference_steps=steps, seed=seed, humanize=humanize, + unsharp=unsharp, max_resolution=max_resolution, + min_resolution=min_resolution, restore_faces=restore_faces, restore_faces_weight=restore_faces_weight, # Detect the vendor from the pristine original (`img_path`), not the @@ -1030,6 +1058,8 @@ def _process_batch_image( help="Cap long side (px) before diffusion; 0 = native (best quality, like raiw.cc). Raise only on GPU/MPS OOM.", ) @_restore_faces_options +@_min_resolution_option +@_unsharp_option @click.pass_context def cmd_batch( ctx: click.Context, @@ -1044,7 +1074,9 @@ def cmd_batch( hf_token: str | None, inpaint: bool, humanize: float, + unsharp: float, max_resolution: int, + min_resolution: int, restore_faces: bool, restore_faces_weight: float, ) -> None: @@ -1096,7 +1128,9 @@ def cmd_batch( seed=seed, hf_token=hf_token, humanize=humanize, + unsharp=unsharp, max_resolution=max_resolution, + min_resolution=min_resolution, restore_faces=restore_faces, restore_faces_weight=restore_faces_weight, ) diff --git a/src/remove_ai_watermarks/humanizer.py b/src/remove_ai_watermarks/humanizer.py index c835ed7..2298fa7 100644 --- a/src/remove_ai_watermarks/humanizer.py +++ b/src/remove_ai_watermarks/humanizer.py @@ -1,7 +1,9 @@ -"""Analog Humanizer: film grain and chromatic aberration injection. +"""Post-processing filters for the cleaned output. -Simulates analog film imperfections to defeat digital AI perfection -classifiers. Ported from NeuralBleach. +``apply_analog_humanizer`` injects film grain and chromatic aberration to defeat +digital AI-perfection classifiers (ported from NeuralBleach); ``unsharp_mask`` +counters the soft, over-smoothed look that diffusion + face-restoration leave +behind (itself a common "this is AI" tell). """ # cv2/numpy boundary: third-party libs ship no usable element types; relax the @@ -56,3 +58,27 @@ def apply_analog_humanizer(image: NDArray, grain_intensity: float = 4.0, chromat humanized = merged return humanized + + +def unsharp_mask(image: NDArray, amount: float = 0.5, sigma: float = 1.0) -> NDArray: + """Sharpen via unsharp masking: ``out = image + amount * (image - blur(image))``. + + Counters the soft, over-smoothed look of the diffusion + GFPGAN passes, which + reads as an AI tell. ``amount`` 0 = no-op (returns an unchanged copy); ~0.5-0.8 + is a safe range -- higher risks bright edge halos that are their own artifact. + ``sigma`` is the Gaussian radius of the unsharp kernel. + + Args: + image: BGR image as numpy array (uint8). + amount: Sharpening strength (0 = off). + sigma: Gaussian blur sigma for the unsharp kernel. + + Returns: + Sharpened BGR image (uint8). + """ + if amount <= 0.0: + return image.copy() + img_f = image.astype(np.float32) + blurred = cv2.GaussianBlur(img_f, (0, 0), sigmaX=sigma, sigmaY=sigma) + sharpened = cv2.addWeighted(img_f, 1.0 + amount, blurred, -amount, 0.0) + return np.clip(sharpened, 0, 255).astype(np.uint8) diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index b25509f..8f8d068 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/src/remove_ai_watermarks/invisible_engine.py @@ -42,21 +42,32 @@ def is_available() -> bool: return importlib.util.find_spec("diffusers") is not None and importlib.util.find_spec("torch") is not None -def _target_size(width: int, height: int, max_resolution: int) -> tuple[int, int] | None: - """Compute the downscaled (width, height) for a long-side cap, or None for native. +def _target_size(width: int, height: int, max_resolution: int, min_resolution: int = 0) -> tuple[int, int] | None: + """Compute the (width, height) to process at, or None for native. - Returns None when no pre-downscale is needed: ``max_resolution <= 0`` (native - resolution, the default that matches the raiw.cc backend -- see issue #10) or - the long side already fits the cap. Otherwise scales the long side down to - ``max_resolution`` preserving aspect ratio (integer-truncated, matching the - PIL ``resize`` call site). Pure function so the native-vs-downscale decision - is unit-testable without loading the diffusion model. + Two opposite long-side adjustments, in precedence order: + + - ``max_resolution`` (cap): if the long side exceeds it, scale DOWN to it + (integer-truncated, matching the PIL ``resize`` call site). 0/negative = no + cap. Set only to bound GPU/MPS memory on very large inputs (issue #10). + - ``min_resolution`` (floor): else if the long side is below it, scale UP to it + (rounded) so SDXL img2img runs near its ~1024 training resolution instead of + degrading on a tiny latent (a 381x512 portrait distorts badly at native). + The output is restored to the original size by the caller, so the floor is a + transparent quality boost. 0 = no floor. Skipped on a ``min > max`` misconfig. + + Returns None when neither applies (native resolution). Pure function so the + resolution decision is unit-testable without loading the diffusion model. """ - if max_resolution > 0 and max(width, height) > max_resolution: - ratio = max_resolution / max(width, height) + long_side = max(width, height) + if max_resolution > 0 and long_side > max_resolution: + ratio = max_resolution / long_side # Clamp the short side to >=1: extreme aspect ratios (e.g. 5000x3 capped # at 1024) would otherwise truncate it to 0 and crash image.resize(). return (max(1, int(width * ratio)), max(1, int(height * ratio))) + if min_resolution > 0 and long_side < min_resolution and (max_resolution <= 0 or min_resolution <= max_resolution): + ratio = min_resolution / long_side + return (max(1, round(width * ratio)), max(1, round(height * ratio))) return None @@ -125,9 +136,11 @@ class InvisibleEngine: seed: int | None = None, humanize: float = 0.0, max_resolution: int = 0, + min_resolution: int = 1024, vendor: str | None = None, restore_faces: bool = False, restore_faces_weight: float = 0.5, + unsharp: float = 0.0, ) -> Path: """Remove invisible watermark from an image. @@ -146,11 +159,20 @@ class InvisibleEngine: absent or no face is detected. restore_faces_weight: GFPGAN fidelity weight (0-1); lower = more GAN regeneration (cleaner watermark scrub), higher = closer to input. + unsharp: Final unsharp-mask sharpening strength (0 = off, default). + Applied last (after face restoration) to counter the soft, + over-smoothed look of the diffusion/GFPGAN passes; ~0.5-0.8 is a + safe range, higher risks edge halos. max_resolution: Cap the long side (px) before diffusion. 0 (default) - = native resolution, no pre-downscale -- matches the hosted - raiw.cc backend. Set a positive value only to bound GPU/MPS - memory on very large inputs (it reintroduces a lossy - downscale->upscale round-trip). + = no cap. Set a positive value only to bound GPU/MPS memory on + very large inputs (it reintroduces a lossy downscale->upscale + round-trip). + min_resolution: Upscale the long side UP to this (px) before diffusion + when the input is smaller, so SDXL runs near its ~1024 training + resolution (small inputs degrade/distort badly at native). 1024 + (default) = on; 0 = off. The output is restored to the original + input size, so this is a transparent quality boost; it adds time + and memory on small inputs. Ignored on a min > max misconfig. Returns: Path to the cleaned image. @@ -159,25 +181,26 @@ class InvisibleEngine: from PIL import Image, ImageOps - # Process at native resolution by default (max_resolution=0). The hosted - # raiw.cc backend (fal fast-sdxl) does NO pre-downscale either, and at - # strength ~0.05 SDXL img2img does not need the input shrunk to ~1024 -- - # the old forced downscale->upscale round-trip was the main quality loss - # (see issue #10). A positive max_resolution caps the long side only to - # bound GPU/MPS memory on very large inputs. + # Resolution policy: a max_resolution cap (0 = none) bounds memory on huge + # inputs, and a min_resolution floor (1024 = default) upscales tiny inputs so + # SDXL img2img runs near its ~1024 training size instead of distorting on a + # tiny latent (a 381x512 portrait wrecks at native -- issue #36 follow-up). + # The output is restored to orig_size below, so the floor is transparent. image = Image.open(image_path) image = ImageOps.exif_transpose(image) orig_size = image.size # (width, height) - # Optional long-side downscale; native resolution by default (issue #10). - target = _target_size(image.width, image.height, max_resolution) + target = _target_size(image.width, image.height, max_resolution, min_resolution) if target is not None: if self._progress_callback: - self._progress_callback( - f"Downscaling {image.width}x{image.height} " - f"to {target[0]}x{target[1]} " - f"(max-resolution cap {max_resolution}px)..." + upscaling = max(target) > max(image.width, image.height) + reason = ( + f"min-resolution floor {min_resolution}px" + if upscaling + else f"max-resolution cap {max_resolution}px" ) + verb = "Upscaling" if upscaling else "Downscaling" + self._progress_callback(f"{verb} {image.width}x{image.height} to {target[0]}x{target[1]} ({reason})...") image = image.resize(target, Image.Resampling.LANCZOS) # Always persist to a temp file, even without downscaling: WatermarkRemover @@ -250,6 +273,20 @@ class InvisibleEngine: if restore_faces: self._restore_faces(out_path, image, restore_faces_weight) + # Final sharpening, LAST so it crisps the face-restored result too (a + # pre-GFPGAN sharpen would be smoothed back over by the face pass). + if unsharp > 0.0: + import cv2 + + from remove_ai_watermarks import image_io + from remove_ai_watermarks.humanizer import unsharp_mask + + out_cv = image_io.imread(out_path, cv2.IMREAD_COLOR) + if out_cv is not None: + if self._progress_callback: + self._progress_callback(f"Sharpening (unsharp mask: {unsharp})...") + image_io.imwrite(out_path, unsharp_mask(out_cv, amount=unsharp)) + return out_path finally: # _tmp_path is always set above (we persist the image unconditionally). diff --git a/tests/test_humanizer.py b/tests/test_humanizer.py index c868b60..9f52545 100644 --- a/tests/test_humanizer.py +++ b/tests/test_humanizer.py @@ -1,6 +1,6 @@ import numpy as np -from remove_ai_watermarks.humanizer import apply_analog_humanizer +from remove_ai_watermarks.humanizer import apply_analog_humanizer, unsharp_mask def test_humanizer_does_not_modify_original_if_disabled(): @@ -70,3 +70,35 @@ def test_chromatic_shift_does_not_wrap_opposite_edge(): # R (index 2) rolled left -> its right border must stay bright (near 255), # NOT wrap the dark left edge. assert result[:, -shift:, 2].min() > 195 + + +def test_unsharp_disabled_returns_unchanged_copy(): + img = np.full((20, 20, 3), 128, dtype=np.uint8) + img[10, 10] = [100, 150, 200] + result = unsharp_mask(img, amount=0.0) + assert np.array_equal(result, img) + assert result is not img # a fresh copy, never the same array + + +def test_unsharp_overshoots_at_an_edge(): + # A vertical step (left 100, right 150). Unsharp masking overshoots at the + # boundary, pushing pixels above the bright level and below the dark level. + img = np.full((20, 20, 3), 100, dtype=np.uint8) + img[:, 10:] = 150 + result = unsharp_mask(img, amount=1.0, sigma=1.5) + assert int(result.max()) > 150 # bright-side overshoot + assert int(result.min()) < 100 # dark-side undershoot + + +def test_unsharp_preserves_shape_and_dtype(): + img = np.full((15, 25, 3), 120, dtype=np.uint8) + result = unsharp_mask(img, amount=0.6) + assert result.shape == img.shape + assert result.dtype == np.uint8 + + +def test_unsharp_flat_image_is_a_noop(): + # No edges -> blur equals the image -> unsharp cancels to the original. + img = np.full((30, 30, 3), 128, dtype=np.uint8) + result = unsharp_mask(img, amount=0.8, sigma=1.0) + assert np.array_equal(result, img) diff --git a/tests/test_invisible_engine.py b/tests/test_invisible_engine.py index 6f76f8f..c7ad063 100644 --- a/tests/test_invisible_engine.py +++ b/tests/test_invisible_engine.py @@ -71,3 +71,33 @@ class TestTargetSize: # the short side must clamp to 1, never 0. assert _target_size(5000, 3, 1024) == (1024, 1) assert _target_size(3, 5000, 1024) == (1, 1024) + + # ── min_resolution floor (small inputs upscaled so SDXL runs near 1024) ── + + def test_floor_default_off(self): + # min_resolution defaults to 0 -> no upscale, preserving legacy behavior. + assert _target_size(381, 512, 0) is None + + def test_floor_upscales_small_input(self): + # 381x512 portrait, floor 1024 -> long side 512 scaled up to 1024 (x2). + assert _target_size(381, 512, 0, 1024) == (762, 1024) + # Landscape: width is the long side. + assert _target_size(512, 381, 0, 1024) == (1024, 762) + + def test_floor_rounds_short_side(self): + # 333x500, floor 1024: ratio 2.048 -> 333*2.048=681.98 rounds to 682. + assert _target_size(333, 500, 0, 1024) == (682, 1024) + + def test_floor_no_op_at_or_above_floor(self): + # Long side already >= floor -> no upscale (and no cap set -> native). + assert _target_size(1024, 768, 0, 1024) is None + assert _target_size(2000, 1000, 0, 1024) is None + + def test_cap_takes_precedence_over_floor(self): + # A huge input with both set: the cap downscales; the floor never fires. + assert _target_size(2000, 1000, 1024, 1024) == (1024, 512) + + def test_floor_skipped_on_min_above_max_misconfig(self): + # min(1024) > max(800) is a misconfig: the floor must not upscale above the + # cap, so it is skipped and the (within-cap) input stays native. + assert _target_size(500, 400, 800, 1024) is None