feat(invisible): sliding-window tiled diffusion for large inputs (--tile)

Add a lossless alternative to the --max-resolution downscale for large
images that OOM on MPS/GPU: regenerate in overlapping, feather-blended
tiles at native resolution.

- noai/tiling.py: pure plan_tiles (uniform tiles, last flush to edge) +
  feather_weights (strictly-positive separable taper -> partition-of-unity
  blend) + run_tiled (per-tile generate callable, decoupled from the
  pipeline). Unit-tested without the model.
- WatermarkRemover.remove_watermark: refactor _generate into _generate_one
  + a tiled branch that engages only when --tile is set and the long side
  exceeds tile_size (ControlNet canny is rebuilt per tile).
- Thread tile/tile_size/tile_overlap through InvisibleEngine and the
  invisible/all/batch CLI commands via a shared _tile_options decorator.

Verified end-to-end on the real SDXL pipeline (forced 2x2 tiling on a
1024px sample, MPS): non-degenerate output, no gross seam at tile borders.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

src/remove_ai_watermarks/cli.py

@@ -189,6 +189,33 @@ _adaptive_polish_option = click.option(
     "no-op there. Pass --no-adaptive-polish to disable. Independent of --unsharp/--humanize.",
 )
 
+
+# Tiled-diffusion knobs, shared by the diffusion commands (invisible/all/batch).
+# Tiling is the lossless alternative to --max-resolution for large inputs that OOM
+# on MPS/GPU: process at native resolution in overlapping, feather-blended tiles.
+def _tile_options(f: Any) -> Any:
+    """Apply the --tile / --tile-size / --tile-overlap options to a command."""
+    f = click.option(
+        "--tile-overlap",
+        type=int,
+        default=128,
+        help="Overlap between adjacent tiles in px (feather-blended, no seam). Default 128.",
+    )(f)
+    f = click.option(
+        "--tile-size",
+        type=int,
+        default=1024,
+        help="Tile dimension in px for --tile (SDXL's training size). Default 1024.",
+    )(f)
+    return click.option(
+        "--tile/--no-tile",
+        default=False,
+        help="Process large images in overlapping tiles instead of one forward pass -- the lossless "
+        "alternative to --max-resolution for inputs that OOM on MPS/GPU. Engages only when the long "
+        "side exceeds --tile-size; pair with --max-resolution 0 (default) to keep native resolution. Default off.",
+    )(f)
+
+
 # HuggingFace model + CFG knobs, shared by the diffusion commands (invisible/all/batch)
 # so the surface stays identical across them.
 _model_option = click.option(
@@ -668,6 +695,7 @@ def cmd_erase(
 @_guidance_scale_option
 @_auto_option
 @_adaptive_polish_option
+@_tile_options
 @click.pass_context
 def cmd_invisible(
     ctx: click.Context,
@@ -689,6 +717,9 @@ def cmd_invisible(
     guidance_scale: float | None,
     auto: bool,
     adaptive_polish: bool,
+    tile: bool,
+    tile_size: int,
+    tile_overlap: int,
 ) -> None:
     """Remove invisible AI watermarks (SynthID, StableSignature, TreeRing).
 
@@ -747,6 +778,9 @@ def cmd_invisible(
         min_resolution=min_resolution,
         upscaler=upscaler,
         vendor=vendor,
+        tile=tile,
+        tile_size=tile_size,
+        tile_overlap=tile_overlap,
     )
     elapsed = time.monotonic() - t0
 
@@ -917,6 +951,7 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo
 @_guidance_scale_option
 @_auto_option
 @_adaptive_polish_option
+@_tile_options
 @click.pass_context
 def cmd_all(
     ctx: click.Context,
@@ -940,6 +975,9 @@ def cmd_all(
     guidance_scale: float | None,
     auto: bool,
     adaptive_polish: bool,
+    tile: bool,
+    tile_size: int,
+    tile_overlap: int,
 ) -> None:
     """Remove ALL watermarks: visible + invisible + metadata.
 
@@ -1044,6 +1082,9 @@ def cmd_all(
                 min_resolution=min_resolution,
                 upscaler=upscaler,
                 vendor=vendor,
+                tile=tile,
+                tile_size=tile_size,
+                tile_overlap=tile_overlap,
             )
             console.print("    Invisible watermark removed")
 
@@ -1121,6 +1162,9 @@ def _process_batch_image(
     model: str | None = None,
     guidance_scale: float | None = None,
     adaptive_polish: bool = False,
+    tile: bool = False,
+    tile_size: int = 1024,
+    tile_overlap: int = 128,
 ) -> None:
     """Process a single image for batch mode.
 
@@ -1179,6 +1223,9 @@ def _process_batch_image(
                 max_resolution=max_resolution,
                 min_resolution=min_resolution,
                 upscaler=upscaler,
+                tile=tile,
+                tile_size=tile_size,
+                tile_overlap=tile_overlap,
                 # Detect the vendor from the pristine original (`img_path`), not the
                 # visible-processed `out_path` whose C2PA is already gone.
                 vendor=vendor_for_strength(img_path),
@@ -1238,6 +1285,7 @@ def _process_batch_image(
 @_guidance_scale_option
 @_auto_option
 @_adaptive_polish_option
+@_tile_options
 @click.pass_context
 def cmd_batch(
     ctx: click.Context,
@@ -1261,6 +1309,9 @@ def cmd_batch(
     guidance_scale: float | None,
     auto: bool,
     adaptive_polish: bool,
+    tile: bool,
+    tile_size: int,
+    tile_overlap: int,
 ) -> None:
     """Process all images in a directory."""
     _banner()
@@ -1321,6 +1372,9 @@ def cmd_batch(
                     model=model,
                     guidance_scale=guidance_scale,
                     adaptive_polish=adaptive_polish,
+                    tile=tile,
+                    tile_size=tile_size,
+                    tile_overlap=tile_overlap,
                 )
                 processed += 1
 

src/remove_ai_watermarks/invisible_engine.py

@@ -170,6 +170,9 @@ class InvisibleEngine:
         unsharp: float = 0.0,
         adaptive_polish: bool = False,
         upscaler: str = "lanczos",
+        tile: bool = False,
+        tile_size: int = 1024,
+        tile_overlap: int = 128,
     ) -> Path:
         """Remove invisible watermark from an image.
 
@@ -205,6 +208,13 @@ class InvisibleEngine:
                 via the ``esrgan`` extra). Only applies when UPscaling (the floor
                 case); a ``max_resolution`` downscale always uses Lanczos. Falls back
                 to Lanczos if the extra is absent.
+            tile: Process the diffusion pass in overlapping tiles instead of one
+                forward pass -- the lossless alternative to ``max_resolution`` for
+                large inputs that OOM on MPS/GPU. Engages only when the long side
+                exceeds ``tile_size``. Pair with ``max_resolution=0`` (the default)
+                so the input keeps its native resolution.
+            tile_size: Tile dimension in px (default 1024).
+            tile_overlap: Overlap between adjacent tiles in px (default 128).
 
         Returns:
             Path to the cleaned image.
@@ -261,6 +271,9 @@ class InvisibleEngine:
                 guidance_scale=guidance_scale,
                 seed=seed,
                 vendor=vendor,
+                tile=tile,
+                tile_size=tile_size,
+                tile_overlap=tile_overlap,
             )
 
             # Post-processing chain: decode the diffusion output ONCE, apply the

src/remove_ai_watermarks/noai/tiling.py

@@ -0,0 +1,142 @@
+"""Sliding-window tiled diffusion for large images.
+
+The img2img / ControlNet pipeline denoises the WHOLE image in one forward pass,
+so it OOMs on MPS/GPU above ~2K (issue #10). Tiling splits the image into
+overlapping tiles -- each kept near SDXL's ~1024 training size -- regenerates
+each tile independently, and feather-blends the overlaps. The result is processed
+at NATIVE resolution with no seam: the lossless alternative to the
+``--max-resolution`` downscale (which trades quality for a smaller forward pass).
+
+The geometry (``plan_tiles``) and the blend weighting (``feather_weights``) are
+pure functions, unit-tested without the diffusion model. ``run_tiled`` is the
+orchestration loop; it takes a ``generate_tile`` callable (one img2img/ControlNet
+pass on a single PIL tile) so it stays decoupled from the pipeline internals.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, NamedTuple
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from numpy.typing import NDArray
+    from PIL import Image as PILImage
+
+logger = logging.getLogger(__name__)
+
+# Strictly-positive floor for the accumulated blend weights so a region covered
+# by a single feathered tile edge (an image corner, no neighbour to blend with)
+# never divides by zero.
+_WEIGHT_EPS = 1e-3
+
+
+class Tile(NamedTuple):
+    """A tile crop box in the source image: top-left ``(x, y)`` + ``width``/``height``."""
+
+    x: int
+    y: int
+    width: int
+    height: int
+
+
+def _axis_positions(length: int, tile: int, overlap: int) -> list[int]:
+    """Tile start offsets along one axis, last tile flush to the far edge.
+
+    Every interior tile is exactly ``tile`` long; the final tile is pulled back
+    to ``length - tile`` so it ends exactly at the edge (it simply overlaps its
+    predecessor a little more). Keeping all tiles the same size is what lets the
+    diffusion pass run at SDXL's preferred dimension on every tile.
+    """
+    if tile <= 0:
+        raise ValueError(f"tile must be positive, got {tile}")
+    if length <= tile:
+        return [0]
+    # Guarantee forward progress even on a pathological overlap >= tile.
+    overlap = min(max(overlap, 0), tile - 1)
+    step = tile - overlap
+    positions = list(range(0, length - tile + 1, step))
+    last = length - tile
+    if positions[-1] != last:
+        positions.append(last)
+    return positions
+
+
+def plan_tiles(width: int, height: int, tile_size: int, overlap: int) -> list[Tile]:
+    """Lay out a grid of overlapping tiles covering ``width`` x ``height``.
+
+    All tiles are ``min(tile_size, width)`` x ``min(tile_size, height)`` (uniform
+    size; the image itself when it fits in one tile). Returned in row-major order.
+    """
+    xs = _axis_positions(width, tile_size, overlap)
+    ys = _axis_positions(height, tile_size, overlap)
+    tile_w = min(tile_size, width)
+    tile_h = min(tile_size, height)
+    return [Tile(x, y, tile_w, tile_h) for y in ys for x in xs]
+
+
+def feather_weights(width: int, height: int, overlap: int) -> NDArray[Any]:
+    """A 2D blend window: ~1 in the interior, ramping down toward each edge.
+
+    Separable linear taper over ``overlap`` pixels from every edge (capped at
+    half the tile so short tiles still taper symmetrically). Strictly positive
+    everywhere, so the normalised blend is well-defined even at an image corner
+    that only one tile covers.
+    """
+    import numpy as np
+
+    def ramp(n: int) -> NDArray[Any]:
+        w = np.ones(n, dtype=np.float32)
+        if overlap > 0 and n > 1:
+            ramp_len = min(overlap, max(1, n // 2))
+            taper = (np.arange(ramp_len, dtype=np.float32) + 1.0) / (ramp_len + 1.0)
+            w[:ramp_len] = taper
+            w[n - ramp_len :] = taper[::-1]
+        return w
+
+    weights = np.outer(ramp(height), ramp(width))
+    np.maximum(weights, _WEIGHT_EPS, out=weights)
+    return weights
+
+
+def run_tiled(
+    generate_tile: Callable[[PILImage.Image], PILImage.Image],
+    image: PILImage.Image,
+    tile_size: int,
+    overlap: int,
+    set_progress: Callable[[str], None] | None = None,
+) -> PILImage.Image:
+    """Tile ``image``, run ``generate_tile`` per tile, and feather-blend the result.
+
+    ``generate_tile`` runs one diffusion pass on a single RGB PIL tile and returns
+    the regenerated tile (the ControlNet control image is built per tile inside it,
+    so each tile gets its own edge map). A pass that rounds dimensions to the latent
+    grid is resized back to the exact tile size before blending.
+    """
+    import numpy as np
+    from PIL import Image
+
+    width, height = image.size
+    tiles = plan_tiles(width, height, tile_size, overlap)
+    accum = np.zeros((height, width, 3), dtype=np.float32)
+    weight_sum = np.zeros((height, width, 1), dtype=np.float32)
+
+    # All tiles share one size (plan_tiles is uniform), so the feather window is
+    # loop-invariant -- compute it once.
+    weights = feather_weights(tiles[0].width, tiles[0].height, overlap)[:, :, None]
+
+    total = len(tiles)
+    for index, tile in enumerate(tiles, start=1):
+        if set_progress is not None:
+            set_progress(f"Tiled diffusion: tile {index}/{total} at ({tile.x},{tile.y}) {tile.width}x{tile.height}...")
+        crop = image.crop((tile.x, tile.y, tile.x + tile.width, tile.y + tile.height))
+        result = generate_tile(crop)
+        if result.size != (tile.width, tile.height):
+            result = result.resize((tile.width, tile.height), Image.Resampling.LANCZOS)
+        arr = np.asarray(result.convert("RGB"), dtype=np.float32)
+        accum[tile.y : tile.y + tile.height, tile.x : tile.x + tile.width] += arr * weights
+        weight_sum[tile.y : tile.y + tile.height, tile.x : tile.x + tile.width] += weights
+
+    blended = accum / np.maximum(weight_sum, _WEIGHT_EPS)
+    return Image.fromarray(np.clip(blended, 0, 255).astype(np.uint8))

src/remove_ai_watermarks/noai/watermark_remover.py

@@ -485,6 +485,9 @@ class WatermarkRemover:
         guidance_scale: float | None = None,
         seed: int | None = None,
         vendor: str | None = None,
+        tile: bool = False,
+        tile_size: int = 1024,
+        tile_overlap: int = 128,
     ) -> Path:
         """Remove watermark from an image using regeneration attack.
 
@@ -501,6 +504,13 @@ class WatermarkRemover:
                 input with ``watermark_profiles.vendor_for_strength`` before processing
                 strips the metadata; the caller passes it down so display and execution
                 agree.
+            tile: Process the image in overlapping tiles instead of one forward pass.
+                The lossless alternative to a ``--max-resolution`` downscale for large
+                inputs that OOM on MPS/GPU (issue #10). Only engages when the long side
+                exceeds ``tile_size``; smaller images run a single pass unchanged.
+            tile_size: Tile dimension in px (default 1024, SDXL's training size).
+            tile_overlap: Overlap between adjacent tiles in px (default 128), feather-
+                blended so there is no visible seam.
 
         Returns:
             Path to the cleaned image.
@@ -541,10 +551,19 @@ class WatermarkRemover:
 
         _total_start = time.monotonic()
 
-        def _generate() -> Image.Image:
+        def _generate_one(img: Image.Image) -> Image.Image:
             if self.model_profile == "controlnet":
-                return self._run_controlnet(init_image, strength, num_inference_steps, guidance_scale, generator)
-            return self._run_img2img(init_image, strength, num_inference_steps, guidance_scale, generator)
+                return self._run_controlnet(img, strength, num_inference_steps, guidance_scale, generator)
+            return self._run_img2img(img, strength, num_inference_steps, guidance_scale, generator)
+
+        def _generate() -> Image.Image:
+            # Tile only when asked AND the image is larger than one tile; otherwise a
+            # single full-image pass (tiling a sub-tile image is pure overhead).
+            if tile and max(init_image.size) > tile_size:
+                from remove_ai_watermarks.noai.tiling import run_tiled
+
+                return run_tiled(_generate_one, init_image, tile_size, tile_overlap, self._set_progress)
+            return _generate_one(init_image)
 
         cleaned_image = _generate()