diff --git a/CLAUDE.md b/CLAUDE.md index c1e2218..c70e277 100644 --- a/CLAUDE.md +++ b/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 -o ` — 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 -o ` — 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 -o ` — 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 -o ` — 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 --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 --region x,y,w,h -o ` — 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 ` — provenance verdict (platform + watermark inventory + confidence); `--json` for machine output, `--no-visible` to skip the cv2 sparkle detector diff --git a/README.md b/README.md index 22afa11..27b8741 100644 --- a/README.md +++ b/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") diff --git a/docs/known-limitations.md b/docs/known-limitations.md index 8f070e8..7a9b502 100644 --- a/docs/known-limitations.md +++ b/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) diff --git a/docs/module-internals.md b/docs/module-internals.md index e5a905f..01c7f27 100644 --- a/docs/module-internals.md +++ b/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 `_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 `_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}`. diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index b66dbf1..0583002 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/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 diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index 9ac7da2..17b3c04 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/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 diff --git a/src/remove_ai_watermarks/noai/watermark_remover.py b/src/remove_ai_watermarks/noai/watermark_remover.py index 95a2c43..9c417be 100644 --- a/src/remove_ai_watermarks/noai/watermark_remover.py +++ b/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 diff --git a/tests/test_cli.py b/tests/test_cli.py index 63d7315..50750f3 100644 --- a/tests/test_cli.py +++ b/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 diff --git a/tests/test_cpu_offload.py b/tests/test_cpu_offload.py new file mode 100644 index 0000000..f2c6ffd --- /dev/null +++ b/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