From 697e11f2ab7828871600863174d58dae95331d59 Mon Sep 17 00:00:00 2001 From: Diego Date: Sat, 5 Sep 2026 21:30:00 +0200 Subject: [PATCH] Report unsupported global CPU offload profiles (#104) Handle unresolved auto preload through its Qwen fallback and avoid a false warning. --- docs/known-limitations.md | 13 +- docs/module-internals.md | 16 ++- .../_internal/watermark_profiles.py | 23 ++++ .../_internal/watermark_remover.py | 20 +++ src/remove_ai_watermarks/cli.py | 6 +- src/remove_ai_watermarks/invisible_engine.py | 7 +- tests/test_cpu_offload.py | 128 +++++++++++++++++- 7 files changed, 196 insertions(+), 17 deletions(-) diff --git a/docs/known-limitations.md b/docs/known-limitations.md index 2790904..83b9d1a 100644 --- a/docs/known-limitations.md +++ b/docs/known-limitations.md @@ -312,9 +312,16 @@ global Qwen stage is tiled; the face stage runs after tile blending. ### CPU offload trades speed for VRAM -`--cpu-offload` forces both stacks of the two-stage profile out of automatic -device residency, streaming weights instead of pinning them. It reduces CUDA -memory pressure at the cost of speed. +`--cpu-offload` forces the shared face stack out of automatic device residency on +every profile, streaming weights instead of pinning them. It reduces CUDA memory +pressure at the cost of speed. + +It reaches the global stack of `qwen-zimage` only. That stack is each +profile's own, and `chroma-zimage` and `sdxl-zimage` load theirs with a plain +`.to(device)`. On a card below the face-stage residency floor, where the face +stack was already offloading, the flag changes nothing at all for those two: the +run behaves as if it were absent. The library says so on the way past, as a +warning ahead of the model load. Residency is otherwise chosen from the card's total VRAM. On a card large enough to hold a stack, offloading is pure waste: DiffSynth drops weights to the meta diff --git a/docs/module-internals.md b/docs/module-internals.md index 2b0a950..3095459 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -1552,8 +1552,20 @@ Regression coverage: CPU offload is enabled only when requested. Nothing calls Diffusers' `enable_model_cpu_offload` any more -- that belonged to the deleted single-stage -profiles. `--cpu-offload` now forces **both** stacks of the two-stage profiles out -of automatic device residency. +profiles. `--cpu-offload` sets both residency fields, and they do not travel the +same distance: the face field is read by the shared base and so applies to every +profile, while the global field is read by `QwenZImagePipeline` alone. +`ChromaZImagePipeline` and `SdxlZImagePipeline` load their global stack with a +plain `.to(device)` and never consult it. + +`GLOBAL_OFFLOAD_PROFILES` in `watermark_profiles.py` is the list, and +`global_offload_supported` the question; `WatermarkRemover._warn_if_global_offload_unsupported` +warns ahead of the model load rather than at construction, because `auto` picks +its engine per-image. An explicit `auto` preload uses the qwen-zimage fallback and +therefore keeps global offload active before that resolution. Below +`RESIDENT_FACE_MODEL_MIN_VRAM_GIB` the flag is a no-op for the other two profiles -- +the face stack was offloading anyway and the global one does not ask -- which is +what the warning exists to say before a download starts. Residency is otherwise chosen from the card's total VRAM, once per stack: `resolve_global_model_residency` gates the mandatory Qwen stack at diff --git a/src/remove_ai_watermarks/_internal/watermark_profiles.py b/src/remove_ai_watermarks/_internal/watermark_profiles.py index ced22e1..4090c31 100644 --- a/src/remove_ai_watermarks/_internal/watermark_profiles.py +++ b/src/remove_ai_watermarks/_internal/watermark_profiles.py @@ -61,6 +61,16 @@ PROFILE_ADAPTIVE_POLISH = { AUTO_PROFILE: False, } +# Which profiles let ``--cpu-offload`` reach their GLOBAL stack. The face stage +# honours the flag everywhere, because residency lives on the shared base, but the +# global stack is each profile's own: only qwen-zimage streams it, through +# ``_qwen_vram_config``. chroma-zimage and sdxl-zimage load theirs with a plain +# ``.to(device)``, so on a card below the face-stage residency floor the flag has +# nothing left to change and the run behaves as if it were absent. Data rather than +# a subclass attribute, because the caller that must warn cannot import a profile +# module without pulling in Diffusers. +GLOBAL_OFFLOAD_PROFILES = frozenset({QWEN_ZIMAGE_PROFILE}) + SDXL_LIGHTNING_MODEL_ID = "ByteDance/SDXL-Lightning" SDXL_LIGHTNING_PATTERN = "sdxl_lightning_4step_lora.safetensors" @@ -221,6 +231,19 @@ def resolve_adaptive_polish(adaptive_polish: bool | None, pipeline: str) -> bool return PROFILE_ADAPTIVE_POLISH.get(normalize_profile(pipeline), True) +def global_offload_supported(pipeline: str) -> bool: + """Whether ``--cpu-offload`` reaches this profile's global stack. + + Per-image ``auto`` runs resolve before loading a stack. An explicit preload is + the only unresolved ``auto`` load, and that path deliberately uses qwen-zimage + as its fallback, so it has the same offload support as qwen-zimage. + """ + profile = normalize_profile(pipeline) + if profile == AUTO_PROFILE: + profile = resolve_auto_profile(None) + return profile in GLOBAL_OFFLOAD_PROFILES + + def strength_default_help() -> str: """Describe the live default policy without duplicating its values.""" return ( diff --git a/src/remove_ai_watermarks/_internal/watermark_remover.py b/src/remove_ai_watermarks/_internal/watermark_remover.py index eb0c741..c2cd414 100644 --- a/src/remove_ai_watermarks/_internal/watermark_remover.py +++ b/src/remove_ai_watermarks/_internal/watermark_remover.py @@ -19,6 +19,7 @@ from remove_ai_watermarks._internal.watermark_profiles import ( QWEN_ZIMAGE_PROFILE, REMOVAL_MODULES, SDXL_ZIMAGE_PROFILE, + global_offload_supported, normalize_profile, resolve_auto_profile, resolve_seed, @@ -153,8 +154,27 @@ class WatermarkRemover: """Materialize the selected model stack before the first request.""" self._load_qwen_zimage_pipeline().preload(global_only=global_only) + def _warn_if_global_offload_unsupported(self) -> None: + """Say so when --cpu-offload cannot reach the stack it was asked for. + + Warned here rather than at construction because ``auto`` picks its engine + per-image. A normal run reaches this point with a concrete profile; preload + reaches it as unresolved ``auto`` and builds its qwen-zimage fallback. + Either way this remains ahead of the model load the user is waiting on. + """ + if not self.cpu_offload or global_offload_supported(self.model_profile): + return + logger.warning( + "--cpu-offload does not reach the global stack of the '%s' pipeline: only " + "qwen-zimage streams its global model between CUDA calls. The face stage " + "still honours the flag, so below the face-stage residency floor this run " + "behaves as if --cpu-offload were absent.", + self.model_profile, + ) + def _load_qwen_zimage_pipeline(self) -> Any: if self._qwen_zimage_pipeline is None: + self._warn_if_global_offload_unsupported() if self.model_profile == SDXL_ZIMAGE_PROFILE: from remove_ai_watermarks._internal.sdxl_zimage_pipeline import ( SdxlZImagePipeline as _Pipeline, diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index a169ed0..031f526 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -337,8 +337,10 @@ _cpu_offload_option = click.option( 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." + "whole pipeline in VRAM, at the cost of speed. Forces the face stack to " + "offload instead of using automatic residency, on every profile. It reaches " + "the global stack of qwen-zimage only: chroma-zimage and sdxl-zimage keep " + "theirs resident either way." ), ) diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index d50703b..4b183df 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/src/remove_ai_watermarks/invisible_engine.py @@ -170,9 +170,10 @@ class InvisibleEngine: controlnet_conditioning_scale: Canny ControlNet structure-preservation strength on the global stage (not used by chroma-zimage). 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. + of keeping the whole pipeline in VRAM, at the cost of speed. Forces + the face stack to offload instead of using automatic residency, on + every profile; reaches the global stack of qwen-zimage only. + CUDA only. """ from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover diff --git a/tests/test_cpu_offload.py b/tests/test_cpu_offload.py index b9fc71a..865522a 100644 --- a/tests/test_cpu_offload.py +++ b/tests/test_cpu_offload.py @@ -1,19 +1,133 @@ -"""Unit tests for the --cpu-offload device-placement branch. +"""Unit tests for how far --cpu-offload travels, per profile. -``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. +The flag sets two residency fields. The face one is read by the shared base and +reaches every profile; the global one is read by ``QwenZImagePipeline`` alone, so +on ``chroma-zimage`` and ``sdxl-zimage`` it never touches the stack the user is +trying to fit. ``WatermarkRemover`` says so before the model load. + +Both halves are exercised with an uninitialized remover, so the core CI matrix +needs no diffusion dependency, model download, or GPU. """ from __future__ import annotations +import logging + +import pytest + +from remove_ai_watermarks._internal.watermark_profiles import ( + AUTO_PROFILE, + CHROMA_ZIMAGE_PROFILE, + GLOBAL_OFFLOAD_PROFILES, + PROFILE_CHOICES, + QWEN_ZIMAGE_PROFILE, + SDXL_ZIMAGE_PROFILE, + global_offload_supported, +) from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover +OFFLOAD_SUPPORT_CASES = ( + (QWEN_ZIMAGE_PROFILE, True), + (CHROMA_ZIMAGE_PROFILE, False), + (SDXL_ZIMAGE_PROFILE, False), + # Unresolved auto appears only in preload, which builds qwen-zimage. + (AUTO_PROFILE, True), +) -def _remover(device: str, cpu_offload: bool) -> WatermarkRemover: + +def _remover(profile: str, cpu_offload: bool) -> WatermarkRemover: remover = WatermarkRemover.__new__(WatermarkRemover) - remover.device = device + remover.model_profile = profile remover.cpu_offload = cpu_offload + remover.device = "cuda" + remover.torch_dtype = None + remover.hf_token = None + remover.controlnet_conditioning_scale = 1.0 remover._progress_callback = None + remover._qwen_zimage_pipeline = None return remover + + +class TestGlobalOffloadSupport: + @pytest.mark.parametrize( + ("profile", "expected"), + OFFLOAD_SUPPORT_CASES, + ) + def test_loaded_profile_or_preload_fallback_reports_support(self, profile: str, expected: bool) -> None: + assert global_offload_supported(profile) is expected + + def test_cases_cover_every_profile(self) -> None: + assert {profile for profile, _expected in OFFLOAD_SUPPORT_CASES} == set(PROFILE_CHOICES) + + def test_the_underscore_spelling_resolves(self) -> None: + assert global_offload_supported("qwen_zimage") is True + + def test_the_set_names_a_profile_that_exists(self) -> None: + # A renamed profile must not leave the flag silently unsupported everywhere. + assert set(PROFILE_CHOICES) >= GLOBAL_OFFLOAD_PROFILES + + +class TestUnsupportedProfileWarns: + @pytest.mark.parametrize("profile", [CHROMA_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE]) + def test_it_names_the_profile_and_the_flag(self, profile: str, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING): + _remover(profile, cpu_offload=True)._warn_if_global_offload_unsupported() + assert len(caplog.records) == 1 + message = caplog.records[0].getMessage() + assert "--cpu-offload" in message + assert profile in message + + def test_qwen_zimage_is_quiet(self, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING): + _remover(QWEN_ZIMAGE_PROFILE, cpu_offload=True)._warn_if_global_offload_unsupported() + assert caplog.records == [] + + def test_a_run_that_did_not_ask_for_offload_is_quiet(self, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING): + _remover(CHROMA_ZIMAGE_PROFILE, cpu_offload=False)._warn_if_global_offload_unsupported() + assert caplog.records == [] + + def test_preloaded_auto_uses_qwen_offload_without_warning( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ) -> None: + from remove_ai_watermarks._internal import qwen_zimage_pipeline + + captured: dict[str, object] = {} + + class Recorder: + def __init__(self, **kwargs: object) -> None: + captured.update(kwargs) + + def preload(self, *, global_only: bool = False) -> None: + captured["global_only"] = global_only + + monkeypatch.setattr(qwen_zimage_pipeline, "QwenZImagePipeline", Recorder) + with caplog.at_level(logging.WARNING): + _remover(AUTO_PROFILE, cpu_offload=True).preload(global_only=True) + + assert caplog.records == [] + assert captured["keep_global_models_on_device"] is False + assert captured["global_only"] is True + + def test_the_load_path_asks_before_building_the_stack(self, monkeypatch: pytest.MonkeyPatch) -> None: + # The warning has to reach the user ahead of the download, and it has to see + # the profile auto resolved to, so the seam is the load and not __init__. + remover = _remover(CHROMA_ZIMAGE_PROFILE, cpu_offload=True) + order: list[str] = [] + monkeypatch.setattr( + WatermarkRemover, + "_warn_if_global_offload_unsupported", + lambda self: order.append("warned"), + ) + + def _explode(*_args: object, **_kwargs: object) -> object: + order.append("loaded") + raise ImportError("no diffusion dependency in this environment") + + monkeypatch.setattr( + "remove_ai_watermarks._internal.chroma_zimage_pipeline.ChromaZImagePipeline", + _explode, + ) + with pytest.raises(ImportError): + remover._load_qwen_zimage_pipeline() + assert order == ["warned", "loaded"]