mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-07-20 06:37:20 +02:00
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 <noreply@anthropic.com>
This commit is contained in:
@@ -10,7 +10,7 @@ You are a **principal Python engineer** maintaining a CLI tool and library for r
|
||||
- `uv run remove-ai-watermarks identify <image>` — provenance verdict (platform + watermark inventory + confidence); `--json` for machine output, `--no-visible` to skip the cv2 sparkle detector
|
||||
- `uv run remove-ai-watermarks metadata <image.png> --check` — inspect AI metadata (C2PA, EXIF, PNG chunks)
|
||||
- `uv run remove-ai-watermarks metadata <image.png> --remove -o <out.png>` — strip all AI metadata
|
||||
- `uv run remove-ai-watermarks batch <directory>` — process every supported image in a directory (output defaults to `<directory>_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 <directory>` — process every supported image in a directory (output defaults to `<directory>_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`.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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).
|
||||
|
||||
+33
-1
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user