Merge pull request #64 from salomonMuriel/feature/cuda-cpu-offload

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

CLAUDE.md

@@ -20,7 +20,7 @@ Consequences for contributors (do not drift back into the stock niche just becau
 Per-command exit-code semantics (the no-signal / GPU-missing skip branches), test traps, and regression-guard paths live in `docs/module-internals.md` (section "CLI commands (`cli.py`)") — read it before changing any command's skip/exit behavior. Every single-image command's `source` argument declares `dir_okay=False`: without it `click.Path(exists=True)` accepts a directory, which then reached `open()` and raised `IsADirectoryError` (Tier E, 2026-07-20; `batch`'s `directory` already declared `file_okay=False`). Regression: `tests/test_cli_robustness.py::TestDirectoryInputIsRejected`.
 
 - `uv run remove-ai-watermarks all <image.png> -o <output.png>` — full pipeline (visible + invisible + metadata). Same diffusion knobs as `invisible`, plus the visible-pass `--backend auto|cv2|migan|lama` (default `auto`) and `--sensitivity auto|strict` (default `auto`) for the localize -> fill visible removal (see the `visible` bullet). Skips step 2 (invisible/SynthID) when the `[gpu]` extra is absent or no invisible signal is detectable; see the module doc for the distinct exit codes.
-- `uv run remove-ai-watermarks invisible <image.png> -o <out.png>` — diffusion SynthID removal. **Full knob set** (kept identical across `invisible`/`all`/`batch`): `--strength` (vendor-adaptive default except resolution-adaptive `qwen-zimage`), `--steps` (**interacts with `--strength`** on the diffusers profiles; `watermark_profiles.viable_steps` prevents zero effective steps. `qwen-zimage` instead fixes its Lightning stage at 4 steps), `--guidance-scale`, `--pipeline sdxl|controlnet|qwen|qwen-zimage` (default `controlnet`; `qwen` and `qwen-zimage` are manual opt-ins), `--controlnet-scale`, `--model`, `--device`, `--seed`, `--hf-token`, `--max-resolution`/`--min-resolution`, `--upscaler lanczos|esrgan`, `--humanize`, `--unsharp`, `--adaptive-polish/--no-adaptive-polish`, `--tile/--no-tile` + `--tile-size`/`--tile-overlap`, `--force/--no-force`. ControlNet is the compatibility and cost default, not the highest-fidelity mode. Recommend the CUDA-only `qwen-zimage` profile when output quality, especially face identity, matters more than runtime and cost; it needs the separate extra, uses a fixed Qwen-Image-2512 + Z-Image stack, rejects `--model`, defaults to the oracle-candidate seed 0, and supports tiling only for its global Qwen pass. The full-frame face stage runs once after tile blending. Tiled outputs still need separate oracle certification. `--auto` is deprecated and a no-op that only warns. Skips the diffusion when no invisible signal is detectable; see the module doc.
+- `uv run remove-ai-watermarks invisible <image.png> -o <out.png>` — diffusion SynthID removal. **Full knob set** (kept identical across `invisible`/`all`/`batch`): `--strength` (vendor-adaptive default except resolution-adaptive `qwen-zimage`), `--steps` (**interacts with `--strength`** on the diffusers profiles; `watermark_profiles.viable_steps` prevents zero effective steps. `qwen-zimage` instead fixes its Lightning stage at 4 steps), `--guidance-scale`, `--pipeline sdxl|controlnet|qwen|qwen-zimage` (default `controlnet`; `qwen` and `qwen-zimage` are manual opt-ins), `--controlnet-scale`, `--model`, `--device`, `--seed`, `--hf-token`, `--max-resolution`/`--min-resolution`, `--upscaler lanczos|esrgan`, `--humanize`, `--unsharp`, `--adaptive-polish/--no-adaptive-polish`, `--tile/--no-tile` + `--tile-size`/`--tile-overlap`, `--cpu-offload/--no-cpu-offload`, `--force/--no-force`. `--cpu-offload` trades speed for lower CUDA VRAM use by moving Diffusers model components between CPU and GPU; on `qwen-zimage` it forces the face stack to offload instead of using automatic residency. It has no effect on CPU/MPS. ControlNet is the compatibility and cost default, not the highest-fidelity mode. Recommend the CUDA-only `qwen-zimage` profile when output quality, especially face identity, matters more than runtime and cost; it needs the separate extra, uses a fixed Qwen-Image-2512 + Z-Image stack, rejects `--model`, defaults to the oracle-candidate seed 0, and supports tiling only for its global Qwen pass. The full-frame face stage runs once after tile blending. Tiled outputs still need separate oracle certification. `--auto` is deprecated and a no-op that only warns. Skips the diffusion when no invisible signal is detectable; see the module doc.
 - `uv run remove-ai-watermarks visible <image.png> -o <out.png>` — known-visible-mark removal by **localize -> fill**: each detected mark is localized to a binary full-frame footprint mask, then one shared, swappable fill inpaints that mask. `--backend auto|cv2|migan|lama` (default `auto`) picks the fill: `cv2` (classical inpaint, no deps, the floor), `migan` (MI-GAN ONNX, light, the memory-tight pick where LaMa will not fit), `lama` (big-LaMa ONNX, best quality, heavier, auto-preferred when a learned backend is available); `auto` = LaMa > MI-GAN > cv2, best available. `--mark auto` (default) removes EVERY detected mark in one pass (a Jimeng-basic image carries the top-left "AI生成" pill AND the bottom-right "★ 即梦AI" wordmark) from: Gemini sparkle, Doubao "豆包AI生成", Jimeng "★ 即梦AI", Qwen "千问AI生成", Samsung Galaxy AI "✦ Contenuti generati dall'AI", Baidu "百度 AI生成", LibLibAI wordmark (bottom-center), RunningHub "RunningHub AI生成" (top-left), and the capture-less Jimeng "AI生成" pill (top-left, metadata-gated); `--mark gemini|doubao|jimeng|qwen|kling|samsung|baidu|liblib|runninghub|jimeng_pill` forces one. `--sensitivity auto|strict` (default `auto`) sets how hard a borderline mark is trusted: `auto` relaxes a mark's gate only on same-product evidence (metadata provenance for that vendor, or a confidently detected sibling mark of the same product — clean images stay untouched); `strict` never relaxes. Metadata provenance is read automatically and feeds `auto`. (`assume-ai` was REMOVED in 0.16 — see the registry bullet; a user who can SEE a missed mark should point at it with `erase --region`, or name it with `--mark <text-mark> --no-detect`.) For arbitrary logos/objects use `erase`. When no known mark is detected the command writes no output and exits with the no-visible-mark code instead of re-serving the input; `--no-detect` forces the gemini fallback and proceeds. See the module doc for the routing/exit detail. `--backend` and `--sensitivity` are shared across `visible`/`all`/`batch`.
 - `uv run remove-ai-watermarks erase <image.png> --region x,y,w,h -o <out.png>` — universal region eraser (any logo/object, any position). `--backend cv2` (default, no deps), `--backend migan` (MI-GAN via onnxruntime, extra `migan`; ~28 MB, ~1 GB RAM, near-LaMa), or `--backend lama` (big-LaMa, extra `lama`; best quality but ~4.7 GB RAM); `--region` is repeatable.
 - `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

README.md

@@ -148,6 +148,7 @@ image → encode to latent space (VAE) at native resolution
 ```
 
 - 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. For inputs that run out of GPU/MPS memory at native resolution, `--tile` is the lossless alternative to `--max-resolution`: it regenerates the image in overlapping, feather-blended tiles so there is no downscale and no visible seam. It engages only when the long side exceeds `--tile-size` (default 1024; overlap `--tile-overlap`, default 128); pair it with `--max-resolution 0`. SDXL and ControlNet run their complete pass per tile. `qwen-zimage` tiles only its global Qwen pass, then runs face detection, SAM masking, and Z-Image restoration once on the blended full frame. 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`. The floor upscale uses Lanczos by default; `--upscaler esrgan` (the `esrgan` extra) runs Real-ESRGAN first for sharper detail and falls back to Lanczos if the extra is absent. ESRGAN is a generic photo/texture GAN with no face/glyph prior, so it is best for photo/texture content -- it can degrade faces (the diffusion pass regenerates them, so the final recovers) and thin text; keep Lanczos for text-heavy inputs.
+- On a memory-constrained CUDA card, `--cpu-offload` moves Diffusers model components between CPU and GPU instead of keeping the full pipeline in VRAM. This lowers peak VRAM use at the cost of transfer overhead. With `qwen-zimage`, the flag forces the face stack to offload instead of using its automatic residency decision. It has no effect on CPU or MPS.
 
 > **Default strength is profile-aware (no flag needed).** SDXL and ControlNet read the C2PA issuer and use **OpenAI gpt-image `0.10`**, **Google Gemini `0.15`**, or **unknown source `0.15`**. The same ladder applies to those two profiles. `qwen` has its own measured ladder. `qwen-zimage` instead ports the upstream resolution-adaptive denoise formula: approximately `0.084` at 0.30 MP through `0.154` at 3.70 MP and above, plus a separate face-adaptive `0.05` through `0.28` pass. There is no local SynthID decoder, so if the corresponding provider oracle still reads the signal, raise `--strength` and re-check the output.
 >
@@ -392,6 +393,8 @@ remove-ai-watermarks invisible image.png -o clean.png --humanize 4.0 --unsharp 0
 # Recommended high-quality mode: full Qwen-Image-2512 Canny pass plus
 # SAM-masked Z-Image face repair. It is much slower and more expensive.
 remove-ai-watermarks invisible image.png -o clean.png --pipeline qwen-zimage --force
+# Low-VRAM CUDA mode. Slower because model components move between CPU and GPU.
+remove-ai-watermarks invisible image.png -o clean.png --cpu-offload --force
 # --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); --upscaler esrgan uses Real-ESRGAN for
@@ -428,7 +431,7 @@ remove-ai-watermarks metadata image.png --remove
 remove-ai-watermarks batch ./images/ --mode visible
 
 # Batch accepts the full invisible knob set (--strength/--guidance-scale/--model/
-# --pipeline/...); --adaptive-polish is on by default except qwen-zimage
+# --pipeline/--cpu-offload/...); --adaptive-polish is on by default except qwen-zimage
 # (--adaptive-polish/--no-adaptive-polish overrides the profile default)
 remove-ai-watermarks batch ./images/ --mode all
 ```
@@ -469,7 +472,8 @@ from remove_ai_watermarks.invisible_engine import InvisibleEngine
 
 # pipeline: "controlnet" (default), "sdxl", "qwen", or CUDA-only "qwen-zimage".
 # model_id=None uses the SDXL base; controlnet_conditioning_scale tunes preservation.
-engine = InvisibleEngine(pipeline="controlnet")
+# cpu_offload=True lowers CUDA VRAM use at the cost of transfer overhead.
+engine = InvisibleEngine(pipeline="controlnet", cpu_offload=True)
 
 # Prefer this CUDA-only profile when output fidelity matters more than cost:
 # engine = InvisibleEngine(pipeline="qwen-zimage")

docs/known-limitations.md

@@ -37,6 +37,8 @@ that cannot afford LaMa's ~4.7 GB peak pins `--backend migan` explicitly.
 
 `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). The floor upscale uses Lanczos by default; **`--upscaler esrgan`** (opt-in, the `esrgan` extra) runs Real-ESRGAN first for better detail before the Lanczos resize to the exact target (`upscaler.py` / `InvisibleEngine._esrgan_upscale`, falls back to Lanczos if the extra is absent). `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.
 
+**CUDA model CPU offload (`--cpu-offload`):** Diffusers normally places the complete SDXL, ControlNet, or base Qwen pipeline on the selected CUDA device. The opt-in offload mode instead uses Diffusers model-level CPU offload, keeping one model component on the GPU at a time and moving it back to CPU before the next component runs. This reduces peak VRAM use but adds transfer overhead. The custom `qwen-zimage` runtime already manages its global model placement; the same flag forces its face stack to stay on the offload path instead of becoming resident on a high-memory GPU. CPU and MPS behavior is unchanged.
+
 **Final `--unsharp` post-filter (`humanizer.unsharp_mask`, opt-in, default 0):** applied LAST (after the face-restore 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` is the lossless alternative** -- see the tiled-diffusion subsection below.
 
 ### Tiled diffusion for large inputs (`--tile`, issue #10)

docs/module-internals.md

@@ -410,7 +410,7 @@ Full pipeline (visible + invisible + metadata). Same diffusion knobs as `invisib
 
 ### `invisible`
 
-Diffusion SynthID removal. The `--tile/--no-tile` knob is the *lossless* alternative to a `--max-resolution` downscale for large inputs that OOM on MPS/GPU: it engages only when the long side exceeds `--tile-size` (default 1024); tiles are feather-blended over `--tile-overlap` px (default 128); pair with `--max-resolution 0`. `--adaptive-polish` is a detail-targeted polish that self-gates to a no-op where there is no deficit; it defaults off only on `qwen-zimage`. `--auto` is deprecated and now a no-op that only warns. **No-signal skip (P0#5, roadmap):** before the diffusion runs, the command checks `identify.has_invisible_target(source)` (the `ProvenanceReport.ai_from_metadata` union: C2PA AI-issuer / SynthID proxy, IPTC, AIGC, local gen params, EXIF/xAI, open DWT-DCT / TrustMark — visible marks do NOT count, they are a separate pass). When nothing is locally detectable it does NOT regenerate (that would only degrade a clean image — the dominant paid score-0 cause on no-watermark uploads): it writes NO output, prints guidance that does NOT claim the image is clean (a pixel SynthID is undetectable once its metadata proxy is gone), and exits **`EXIT_NO_INVISIBLE_SIGNAL` (2)** — same value/role as the visible `EXIT_NO_VISIBLE_MARK`. `--force/--no-force` (**default skip = ON**) runs the scrub regardless. The check fails SAFE (a detector exception → run, since leaving a watermark on a paid removal is worse than over-regenerating). Helpers `cli._no_invisible_signal_exit` + `identify.has_invisible_target`; regression-guarded by `tests/test_cli.py::TestInvisibleCommand::{test_invisible_no_signal_skips_and_exits_two,test_invisible_force_runs_scrub_on_no_signal,test_invisible_runs_without_force_when_signal_present}` and `tests/test_identify.py::TestHasInvisibleTargetFailSafe`. **Test trap:** any `invisible`/`all`/`batch` test that exercises the diffusion path on a signal-LESS fixture (e.g. the synthetic `sample_png`) MUST pass `--force`, or the new gate skips step 2 (so `mock_engine.remove_watermark` is never called / `invisible` exits 2).
+Diffusion SynthID removal. The `--tile/--no-tile` knob is the *lossless* alternative to a `--max-resolution` downscale for large inputs that OOM on MPS/GPU: it engages only when the long side exceeds `--tile-size` (default 1024); tiles are feather-blended over `--tile-overlap` px (default 128); pair with `--max-resolution 0`. `--cpu-offload/--no-cpu-offload` trades speed for lower CUDA VRAM use: SDXL, ControlNet, and base Qwen call Diffusers `enable_model_cpu_offload(device="cuda")`, which moves whole model components between CPU and GPU; `qwen-zimage` instead forces its face stack to use the existing offload path rather than automatic high-VRAM residency. The flag has no effect on CPU/MPS and fails loudly if a CUDA Diffusers pipeline lacks the offload method. `--adaptive-polish` is a detail-targeted polish that self-gates to a no-op where there is no deficit; it defaults off only on `qwen-zimage`. `--auto` is deprecated and now a no-op that only warns. **No-signal skip (P0#5, roadmap):** before the diffusion runs, the command checks `identify.has_invisible_target(source)` (the `ProvenanceReport.ai_from_metadata` union: C2PA AI-issuer / SynthID proxy, IPTC, AIGC, local gen params, EXIF/xAI, open DWT-DCT / TrustMark — visible marks do NOT count, they are a separate pass). When nothing is locally detectable it does NOT regenerate (that would only degrade a clean image — the dominant paid score-0 cause on no-watermark uploads): it writes NO output, prints guidance that does NOT claim the image is clean (a pixel SynthID is undetectable once its metadata proxy is gone), and exits **`EXIT_NO_INVISIBLE_SIGNAL` (2)** — same value/role as the visible `EXIT_NO_VISIBLE_MARK`. `--force/--no-force` (**default skip = ON**) runs the scrub regardless. The check fails SAFE (a detector exception → run, since leaving a watermark on a paid removal is worse than over-regenerating). Helpers `cli._no_invisible_signal_exit` + `identify.has_invisible_target`; regression-guarded by `tests/test_cli.py::TestInvisibleCommand::{test_invisible_no_signal_skips_and_exits_two,test_invisible_force_runs_scrub_on_no_signal,test_invisible_runs_without_force_when_signal_present,test_invisible_cpu_offload_flows_to_engine}`, `tests/test_cli.py::TestAllCommand::test_all_cpu_offload_flows_to_engine`, `tests/test_cli.py::TestBatchCommand::test_batch_cpu_offload_flows_to_cached_engine`, and `tests/test_identify.py::TestHasInvisibleTargetFailSafe`. **Test trap:** any `invisible`/`all`/`batch` test that exercises the diffusion path on a signal-LESS fixture (e.g. the synthetic `sample_png`) MUST pass `--force`, or the new gate skips step 2 (so `mock_engine.remove_watermark` is never called / `invisible` exits 2).
 
 ### `visible`
 
@@ -418,4 +418,4 @@ Known-visible-mark removal by **localize -> fill**: each detected mark is locali
 
 ### `batch`
 
-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 **full `invisible` knob set** (`--strength`/`--steps`/`--guidance-scale`/`--pipeline`/`--controlnet-scale`/`--model`/`--device`/`--max-resolution`/`--min-resolution`/`--upscaler`/`--seed`/`--hf-token`/`--humanize`/`--unsharp`/`--adaptive-polish`/`--tile`/`--tile-size`/`--tile-overlap`/`--force`), plus `--backend` for the visible localize -> fill pass. `--adaptive-polish` is ON by default except on `qwen-zimage`; `--auto` is deprecated and a no-op that only warns. **No-signal skip (P0#5):** in invisible/all mode each image runs the same `has_invisible_target` gate — a signal-less image is skipped (no diffusion); in `invisible` mode the input is copied through to the output dir so it stays complete, in `all` mode the visible-removed result is kept and metadata is still stripped. `--force` scrubs every image regardless. One engine cached per pipeline; the polish is resolved once before the loop. **Exit code (`batch` used to always exit 0, hiding failures):** `cmd_batch` raises `SystemExit(1)` when any image errored, OR when a `--mode invisible`/`all` image carried an invisible signal but the GPU extra was absent so its SynthID scrub was skipped — mirroring single `all`, it emits a loud "the invisible watermark was NOT removed on N image(s)" warning and (invisible mode) copies the input through so the output dir stays complete, rather than silently dropping the signal-bearing files that most needed processing. `_process_batch_image` returns that skipped-scrub flag; the loop tallies it. Regression-guarded by `tests/test_cli.py::TestBatchCommand::{test_batch_errors_exit_nonzero, test_batch_invisible_gpu_missing_writes_output_and_exits_nonzero}`.
+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 **full `invisible` knob set** (`--strength`/`--steps`/`--guidance-scale`/`--pipeline`/`--controlnet-scale`/`--model`/`--device`/`--max-resolution`/`--min-resolution`/`--upscaler`/`--seed`/`--hf-token`/`--humanize`/`--unsharp`/`--adaptive-polish`/`--tile`/`--tile-size`/`--tile-overlap`/`--cpu-offload`/`--force`), plus `--backend` for the visible localize -> fill pass. `--adaptive-polish` is ON by default except on `qwen-zimage`; `--auto` is deprecated and a no-op that only warns. **No-signal skip (P0#5):** in invisible/all mode each image runs the same `has_invisible_target` gate — a signal-less image is skipped (no diffusion); in `invisible` mode the input is copied through to the output dir so it stays complete, in `all` mode the visible-removed result is kept and metadata is still stripped. `--force` scrubs every image regardless. One engine cached per pipeline; the polish is resolved once before the loop. **Exit code (`batch` used to always exit 0, hiding failures):** `cmd_batch` raises `SystemExit(1)` when any image errored, OR when a `--mode invisible`/`all` image carried an invisible signal but the GPU extra was absent so its SynthID scrub was skipped — mirroring single `all`, it emits a loud "the invisible watermark was NOT removed on N image(s)" warning and (invisible mode) copies the input through so the output dir stays complete, rather than silently dropping the signal-bearing files that most needed processing. `_process_batch_image` returns that skipped-scrub flag; the loop tallies it. Regression-guarded by `tests/test_cli.py::TestBatchCommand::{test_batch_errors_exit_nonzero, test_batch_invisible_gpu_missing_writes_output_and_exits_nonzero}`.

src/remove_ai_watermarks/cli.py

@@ -312,6 +312,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=(
+        "Offload model components to CPU between CUDA calls instead of keeping the "
+        "whole pipeline in VRAM, at the cost of speed. For qwen-zimage, forces the "
+        "face stack to offload instead of using automatic residency. No effect on "
+        "cpu/mps."
+    ),
+)
 
 
 _visible_backend_option = click.option(
@@ -901,6 +911,7 @@ def cmd_erase(
 @_adaptive_polish_option
 @_tile_options
 @_force_option
+@_cpu_offload_option
 @click.pass_context
 def cmd_invisible(
     ctx: click.Context,
@@ -926,6 +937,7 @@ def cmd_invisible(
     tile_size: int,
     tile_overlap: int,
     force: bool,
+    cpu_offload: bool,
 ) -> None:
     """Remove invisible AI watermarks (SynthID, StableSignature, TreeRing).
 
@@ -969,6 +981,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
@@ -1187,6 +1200,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,
@@ -1214,6 +1228,7 @@ def cmd_all(
     tile_size: int,
     tile_overlap: int,
     force: bool,
+    cpu_offload: bool,
 ) -> None:
     """Remove ALL watermarks: visible + invisible + metadata.
 
@@ -1313,6 +1328,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
@@ -1433,6 +1449,7 @@ class _BatchOptions:
     tile_size: int = 1024
     tile_overlap: int = 128
     force: bool = False
+    cpu_offload: bool = False
 
 
 def _run_batch_invisible(
@@ -1466,6 +1483,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,
@@ -1623,6 +1641,7 @@ def _process_batch_image(
 @_adaptive_polish_option
 @_tile_options
 @_force_option
+@_cpu_offload_option
 @click.pass_context
 def cmd_batch(
     ctx: click.Context,
@@ -1651,6 +1670,7 @@ def cmd_batch(
     tile_size: int,
     tile_overlap: int,
     force: bool,
+    cpu_offload: bool,
 ) -> None:
     """Process all images in a directory."""
     _banner()
@@ -1695,6 +1715,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

@@ -100,6 +100,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.
 
@@ -117,6 +118,10 @@ class InvisibleEngine:
             progress_callback: Optional callback for progress messages.
             controlnet_conditioning_scale: ControlNet structure-preservation
                 strength (controlnet pipeline only).
+            cpu_offload: Offload model components to CPU between CUDA calls instead
+                of keeping the whole pipeline in VRAM, at the cost of speed. For
+                qwen-zimage, force the face stack to offload instead of using automatic
+                residency. CUDA only.
         """
 
         from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
@@ -130,6 +135,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

@@ -387,8 +387,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
+        # Diffusers offloads whole model components between CUDA calls. The custom
+        # qwen-zimage runtime uses the same flag to keep its face stack off VRAM.
+        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".
@@ -484,17 +488,29 @@ class WatermarkRemover:
         trigger the torch-CUDA reinstall+restart. Returns the moved pipeline.
         """
         self._set_progress(f"Moving model to device: {self.device}")
-        try:
-            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...")
-                _reinstall_torch_cuda_and_restart()
-            raise RuntimeError(
-                f"Failed to move model to {self.device} ({exc}). "
-                "Install CUDA-enabled PyTorch manually:\n"
-                f"  pip install torch --index-url {_detect_cuda_index_url()}"
-            ) from exc
+        if self.cpu_offload and self.device == "cuda":
+            enable_cpu_offload = getattr(pipeline, "enable_model_cpu_offload", None)
+            if not callable(enable_cpu_offload):
+                raise RuntimeError(
+                    "CPU offload was requested, but this pipeline does not support enable_model_cpu_offload()."
+                )
+            self._set_progress("Enabling CUDA model CPU offload (low-VRAM mode)...")
+            try:
+                enable_cpu_offload(device=self.device)
+            except (RuntimeError, AssertionError) as exc:
+                raise RuntimeError(f"Failed to enable model CPU offload ({exc}).") from exc
+        else:
+            try:
+                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...")
+                    _reinstall_torch_cuda_and_restart()
+                raise RuntimeError(
+                    f"Failed to move model to {self.device} ({exc}). "
+                    "Install CUDA-enabled PyTorch manually:\n"
+                    f"  pip install torch --index-url {_detect_cuda_index_url()}"
+                ) from exc
 
         if hasattr(pipeline, "enable_xformers_memory_efficient_attention"):
             with contextlib.suppress(Exception):
@@ -638,6 +654,7 @@ class WatermarkRemover:
                 hf_token=self.hf_token,
                 progress_callback=self._progress_callback,
                 controlnet_conditioning_scale=self.controlnet_conditioning_scale,
+                keep_face_models_on_device=False if self.cpu_offload else None,
             )
         return self._qwen_zimage_pipeline
 

tests/test_cli.py

@@ -299,6 +299,22 @@ class TestInvisibleCommand:
         assert output.exists()
         mock_engine.remove_watermark.assert_called_once()
 
+    def test_invisible_cpu_offload_flows_to_engine(self, runner, sample_png, tmp_path):
+        mock_cls, _mock_engine = _mock_invisible_engine()
+        output = tmp_path / "clean.png"
+        with (
+            patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
+            patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
+            patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
+        ):
+            result = runner.invoke(
+                main,
+                ["invisible", str(sample_png), "-o", str(output), "--cpu-offload", "--force"],
+            )
+
+        assert result.exit_code == 0, result.output
+        assert mock_cls.call_args.kwargs["cpu_offload"] is True
+
     def test_invisible_default_output(self, runner, sample_png):
         mock_cls, _mock_engine = _mock_invisible_engine()
         with (
@@ -456,6 +472,22 @@ class TestAllCommand:
         assert result.exit_code == 0, result.output
         assert output.exists()
 
+    def test_all_cpu_offload_flows_to_engine(self, runner, sample_png, tmp_path):
+        mock_cls, _mock_engine = _mock_invisible_engine()
+        output = tmp_path / "clean.png"
+        with (
+            patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
+            patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
+            patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
+        ):
+            result = runner.invoke(
+                main,
+                ["all", str(sample_png), "-o", str(output), "--cpu-offload", "--force"],
+            )
+
+        assert result.exit_code == 0, result.output
+        assert mock_cls.call_args.kwargs["cpu_offload"] is True
+
     def test_all_nonexistent_file(self, runner):
         result = runner.invoke(main, ["all", "/nonexistent/file.png"])
         assert result.exit_code != 0
@@ -718,6 +750,33 @@ class TestBatchCommand:
         assert result.exit_code == 0, result.output
         assert "3 processed" in result.output
 
+    def test_batch_cpu_offload_flows_to_cached_engine(self, runner, tmp_path):
+        input_dir = _make_batch_dir(tmp_path)
+        output_dir = tmp_path / "output"
+        mock_cls, _mock_engine = _mock_invisible_engine()
+        with (
+            patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
+            patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
+            patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True),
+            patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
+        ):
+            result = runner.invoke(
+                main,
+                [
+                    "batch",
+                    str(input_dir),
+                    "-o",
+                    str(output_dir),
+                    "--mode",
+                    "invisible",
+                    "--cpu-offload",
+                    "--force",
+                ],
+            )
+
+        assert result.exit_code == 0, result.output
+        assert mock_cls.call_args.kwargs["cpu_offload"] is True
+
     def test_batch_invisible_skips_no_signal_and_copies_through(self, runner, tmp_path):
         """P0#5: batch invisible mode skips the scrub on signal-less images (no
         --force) and copies the input through, so the output dir is complete with the

tests/test_cpu_offload.py

@@ -0,0 +1,75 @@
+"""Unit tests for the --cpu-offload device-placement branch.
+
+``WatermarkRemover._move_to_device_and_optimize`` chooses between a full
+``pipeline.to("cuda")`` and ``enable_model_cpu_offload()``. The placement
+decision is exercised with a mock pipeline and an uninitialized remover, so the
+core CI matrix needs no diffusion dependency, model download, or GPU.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import Mock
+
+import pytest
+
+from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
+
+
+def _remover(device: str, cpu_offload: bool) -> WatermarkRemover:
+    remover = WatermarkRemover.__new__(WatermarkRemover)
+    remover.device = device
+    remover.cpu_offload = cpu_offload
+    remover._progress_callback = None
+    return remover
+
+
+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(device="cuda")
+        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()
+
+    def test_offload_fails_loudly_when_pipeline_lacks_support(self):
+        remover = _remover("cuda", cpu_offload=True)
+        pipeline = Mock(spec=["to"])
+
+        with pytest.raises(RuntimeError, match="does not support"):
+            remover._move_to_device_and_optimize(pipeline)
+
+        pipeline.to.assert_not_called()
+
+    def test_qwen_zimage_forces_face_stack_offload(self):
+        remover = _remover("cuda", cpu_offload=True)
+        remover.torch_dtype = object()
+        remover.hf_token = None
+        remover.controlnet_conditioning_scale = 1.0
+        remover._qwen_zimage_pipeline = None
+
+        runtime = remover._load_qwen_zimage_pipeline()
+
+        assert runtime.keep_face_models_on_device is False