Add --cpu-offload flag for low-VRAM CUDA cards

The invisible/SynthID diffusion pass loads the whole SDXL fp16 pipeline
into VRAM via `pipeline.to("cuda")`. On an 8 GB card the weights alone
(~7 GB) leave no room for activations, so the run OOMs and there is no
in-tool way to recover short of falling back to CPU (~9 min/image).

Add an opt-in `--cpu-offload` flag (default off) that calls diffusers'
`enable_model_cpu_offload()` instead: submodules are streamed to the GPU
on demand, dropping peak VRAM to roughly the largest single submodule at
the cost of per-step transfers. CUDA-only; a no-op on cpu/mps. Threaded
through `invisible`, `all`, and `batch` to keep the knob set identical
across the three, mirroring the existing `--device`/`--pipeline` options.

Measured on a GTX 1070 Ti (8 GB): `invisible --pipeline sdxl --cpu-offload`
runs the SynthID scrub on-GPU in ~2.5 min vs ~9 min on CPU, where the
default full-VRAM path OOMs.

Test drives the placement decision with a mock pipeline (no model/GPU),
gated on torch so it runs under the `gpu` extra and skips the core CI
matrix, consistent with the model-running test policy.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

src/remove_ai_watermarks/cli.py

@@ -290,6 +290,16 @@ _force_option = click.option(
         "undetectable once its metadata proxy is gone)."
     ),
 )
+_cpu_offload_option = click.option(
+    "--cpu-offload/--no-cpu-offload",
+    default=False,
+    help=(
+        "Stream pipeline submodules to the GPU on demand instead of holding the whole "
+        "fp16 pipeline in VRAM (CUDA only). Lets a low-VRAM card (e.g. 8 GB) run SDXL "
+        "that would otherwise OOM, at the cost of speed. Pair with --pipeline sdxl on "
+        "the tightest cards. No effect on cpu/mps."
+    ),
+)
 
 
 _visible_backend_option = click.option(
@@ -856,6 +866,7 @@ def cmd_erase(
 @_adaptive_polish_option
 @_tile_options
 @_force_option
+@_cpu_offload_option
 @click.pass_context
 def cmd_invisible(
     ctx: click.Context,
@@ -881,6 +892,7 @@ def cmd_invisible(
     tile_size: int,
     tile_overlap: int,
     force: bool,
+    cpu_offload: bool,
 ) -> None:
     """Remove invisible AI watermarks (SynthID, StableSignature, TreeRing).
 
@@ -922,6 +934,7 @@ def cmd_invisible(
         hf_token=hf_token,
         progress_callback=progress_cb,
         controlnet_conditioning_scale=controlnet_scale,
+        cpu_offload=cpu_offload,
     )
 
     # Detect the SynthID vendor from the ORIGINAL (before processing strips C2PA) so the
@@ -1130,6 +1143,7 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo
 @_adaptive_polish_option
 @_tile_options
 @_force_option
+@_cpu_offload_option
 @click.pass_context
 def cmd_all(
     ctx: click.Context,
@@ -1157,6 +1171,7 @@ def cmd_all(
     tile_size: int,
     tile_overlap: int,
     force: bool,
+    cpu_offload: bool,
 ) -> None:
     """Remove ALL watermarks: visible + invisible + metadata.
 
@@ -1254,6 +1269,7 @@ def cmd_all(
                 hf_token=hf_token,
                 progress_callback=progress_cb,
                 controlnet_conditioning_scale=controlnet_scale,
+                cpu_offload=cpu_offload,
             )
 
             # Detect the vendor from the pristine ORIGINAL (`source`); `tmp_path` has
@@ -1372,6 +1388,7 @@ class _BatchOptions:
     tile_size: int = 1024
     tile_overlap: int = 128
     force: bool = False
+    cpu_offload: bool = False
 
 
 def _run_batch_invisible(
@@ -1405,6 +1422,7 @@ def _run_batch_invisible(
                 pipeline=options.pipeline,
                 hf_token=options.hf_token,
                 controlnet_conditioning_scale=options.controlnet_scale,
+                cpu_offload=options.cpu_offload,
             )
         engines[options.pipeline].remove_watermark(
             img_path if mode == "invisible" else out_path,
@@ -1552,6 +1570,7 @@ def _process_batch_image(
 @_adaptive_polish_option
 @_tile_options
 @_force_option
+@_cpu_offload_option
 @click.pass_context
 def cmd_batch(
     ctx: click.Context,
@@ -1580,6 +1599,7 @@ def cmd_batch(
     tile_size: int,
     tile_overlap: int,
     force: bool,
+    cpu_offload: bool,
 ) -> None:
     """Process all images in a directory."""
     _banner()
@@ -1622,6 +1642,7 @@ def cmd_batch(
         tile_size=tile_size,
         tile_overlap=tile_overlap,
         force=force,
+        cpu_offload=cpu_offload,
     )
 
     processed = 0

src/remove_ai_watermarks/invisible_engine.py

@@ -95,6 +95,7 @@ class InvisibleEngine:
         hf_token: str | None = None,
         progress_callback: Callable[[str], None] | None = None,
         controlnet_conditioning_scale: float = 1.0,
+        cpu_offload: bool = False,
     ) -> None:
         """Initialize the invisible watermark removal engine.
 
@@ -110,6 +111,9 @@ class InvisibleEngine:
             progress_callback: Optional callback for progress messages.
             controlnet_conditioning_scale: ControlNet structure-preservation
                 strength (controlnet pipeline only).
+            cpu_offload: Stream pipeline submodules to CUDA on demand instead of
+                holding the whole fp16 pipeline in VRAM. Lets a low-VRAM card (e.g.
+                8 GB) run SDXL that would otherwise OOM, at the cost of speed. CUDA only.
         """
 
         from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
@@ -123,6 +127,7 @@ class InvisibleEngine:
             hf_token=hf_token,
             pipeline=pipeline,
             controlnet_conditioning_scale=controlnet_conditioning_scale,
+            cpu_offload=cpu_offload,
         )
         self._progress_callback = progress_callback
 

src/remove_ai_watermarks/noai/watermark_remover.py

@@ -381,8 +381,12 @@ class WatermarkRemover:
         hf_token: str | None = None,
         pipeline: str = "controlnet",
         controlnet_conditioning_scale: float = 1.0,
+        cpu_offload: bool = False,
     ) -> None:
         self.model_id = model_id or self.DEFAULT_MODEL_ID
+        # Stream pipeline submodules to CUDA on demand instead of holding the whole
+        # fp16 pipeline in VRAM -- lets an 8 GB card run SDXL that would otherwise OOM.
+        self.cpu_offload = cpu_offload
         # The pipeline profile is threaded explicitly (not inferred from model_id):
         # both "sdxl" and "controlnet" use the same SDXL base checkpoint. Normalize so
         # the legacy "default" alias resolves to "sdxl".
@@ -469,7 +473,15 @@ class WatermarkRemover:
         """
         self._set_progress(f"Moving model to device: {self.device}")
         try:
-            pipeline = pipeline.to(self.device)
+            # Low-VRAM CUDA cards (e.g. an 8 GB Pascal card) cannot hold the whole SDXL
+            # fp16 pipeline in VRAM. With --cpu-offload, stream submodules to the GPU on
+            # demand (accelerate hooks) instead of a full .to("cuda"): peak VRAM drops to
+            # roughly the largest single submodule, at the cost of per-step transfers.
+            if self.cpu_offload and self.device == "cuda" and hasattr(pipeline, "enable_model_cpu_offload"):
+                self._set_progress("Enabling CUDA model CPU offload (low-VRAM mode)...")
+                pipeline.enable_model_cpu_offload()
+            else:
+                pipeline = pipeline.to(self.device)
         except (RuntimeError, AssertionError) as exc:
             if self.device == "cuda" and not os.environ.get(_CUDA_FIX_ENV_KEY):
                 self._set_progress("CUDA failed. Reinstalling torch with CUDA support...")

tests/test_cpu_offload.py

@@ -0,0 +1,56 @@
+"""Unit tests for the --cpu-offload device-placement branch (mocked pipeline).
+
+``WatermarkRemover._move_to_device_and_optimize`` chooses between a full
+``pipeline.to("cuda")`` and ``enable_model_cpu_offload()`` (low-VRAM streaming).
+Constructing the remover is cheap -- the diffusion pipeline is lazy and the
+device string is not validated -- so the placement decision is exercised with a
+mock pipeline, no model download or GPU required. Gated on torch (the module
+imports it at top), so it runs under the ``gpu`` extra and skips the core CI
+matrix, matching the model-running test policy.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import Mock
+
+import pytest
+
+pytest.importorskip("torch")
+
+from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
+
+
+def _remover(device: str, cpu_offload: bool) -> WatermarkRemover:
+    return WatermarkRemover(device=device, pipeline="sdxl", cpu_offload=cpu_offload)
+
+
+class TestCpuOffloadPlacement:
+    def test_offload_enabled_on_cuda_streams_instead_of_moving(self):
+        remover = _remover("cuda", cpu_offload=True)
+        pipeline = Mock()
+
+        returned = remover._move_to_device_and_optimize(pipeline)
+
+        pipeline.enable_model_cpu_offload.assert_called_once_with()
+        pipeline.to.assert_not_called()
+        # Offload leaves the pipeline object in place (accelerate hooks handle it).
+        assert returned is pipeline
+
+    def test_no_offload_moves_whole_pipeline_to_cuda(self):
+        remover = _remover("cuda", cpu_offload=False)
+        pipeline = Mock()
+
+        remover._move_to_device_and_optimize(pipeline)
+
+        pipeline.to.assert_called_once_with("cuda")
+        pipeline.enable_model_cpu_offload.assert_not_called()
+
+    def test_offload_flag_ignored_off_cuda(self):
+        # The flag is CUDA-only: on cpu it must still be a plain .to("cpu").
+        remover = _remover("cpu", cpu_offload=True)
+        pipeline = Mock()
+
+        remover._move_to_device_and_optimize(pipeline)
+
+        pipeline.to.assert_called_once_with("cpu")
+        pipeline.enable_model_cpu_offload.assert_not_called()