From 39fdd59f6c2e5c8000de6725536e40477773306e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Salom=C3=B3n=20Muriel?= Date: Thu, 23 Jul 2026 16:06:43 -0500 Subject: [PATCH] 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) --- src/remove_ai_watermarks/cli.py | 21 +++++++ src/remove_ai_watermarks/invisible_engine.py | 5 ++ .../noai/watermark_remover.py | 14 ++++- tests/test_cpu_offload.py | 56 +++++++++++++++++++ 4 files changed, 95 insertions(+), 1 deletion(-) create mode 100644 tests/test_cpu_offload.py diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index f2253d3..ecf5847 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/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 diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index 977c5f6..80bd281 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/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 diff --git a/src/remove_ai_watermarks/noai/watermark_remover.py b/src/remove_ai_watermarks/noai/watermark_remover.py index bdd392f..a2c4856 100644 --- a/src/remove_ai_watermarks/noai/watermark_remover.py +++ b/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...") diff --git a/tests/test_cpu_offload.py b/tests/test_cpu_offload.py new file mode 100644 index 0000000..89d4df7 --- /dev/null +++ b/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()