Merge pull request #117 from dichiara19/fix/sdxl-zimage-global-offload

Stream the sdxl-zimage global stack on cards below 12 GiB

docs/known-limitations.md

@@ -400,12 +400,20 @@ global Qwen stage is tiled; the face stage runs after tile blending.
 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.
+It reaches the global stack of `qwen-zimage` and `sdxl-zimage`. That stack is
+each profile's own, and `chroma-zimage` loads its stack 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 that profile: the run behaves as
+if it were absent. The library says so on the way past, as a warning ahead of the
+model load.
+
+`sdxl-zimage` streams its global stack without being asked on a card below
+`RESIDENT_SDXL_GLOBAL_MIN_VRAM_GIB` (12 GiB). Its UNet and Canny ControlNet run
+together on every step, so the stack does not fit an 8 GiB card even with
+model-level offload. A resident load there does not always fail loudly: under
+the Windows display driver, including WSL2, CUDA can spill the overflow into
+shared system memory, and the run goes on at a fraction of the speed instead of
+raising an out-of-memory error.
 
 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

docs/module-internals.md

@@ -2185,27 +2185,27 @@ Regression coverage:
 
 ### CPU offload
 
-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` 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.
+Nothing calls Diffusers' `enable_model_cpu_offload` any more -- that belonged to
+the deleted single-stage 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`
+and `SdxlZImagePipeline`. `ChromaZImagePipeline` loads its global stack with a
+plain `.to(device)` and never consults 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.
+`RESIDENT_FACE_MODEL_MIN_VRAM_GIB` the flag is a no-op for `chroma-zimage` -- 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
-`RESIDENT_GLOBAL_MODEL_MIN_VRAM_GIB` and `resolve_face_model_residency` gates the
-optional Z-Image stack at `RESIDENT_FACE_MODEL_MIN_VRAM_GIB`.
+`RESIDENT_GLOBAL_MODEL_MIN_VRAM_GIB`, `resolve_sdxl_global_residency` gates the
+SDXL stack at `RESIDENT_SDXL_GLOBAL_MIN_VRAM_GIB`, and `resolve_face_model_residency`
+gates the optional Z-Image stack at `RESIDENT_FACE_MODEL_MIN_VRAM_GIB`.
 
 Below the global floor, `_qwen_vram_config` streams the stack from disk, which is
 what makes a 20B model runnable on a consumer card. At or above it, streaming is
@@ -2222,6 +2222,43 @@ The resident config deliberately passes no `"disk"` value anywhere. DiffSynth la
 `disk_offload` once, from `offload_dtype`, so leaving the sentinel in place while
 pointing every device at CUDA would keep the meta-drop and re-read.
 
+Sequential SDXL offload receives the current CUDA device index explicitly, so
+its hooks use the same GPU as the residency probes, input tensors, and seeded
+generators. Omitting the index makes Diffusers select GPU 0 even after
+`torch.cuda.set_device(1)`. `test_streaming_uses_the_current_cuda_device` in
+`test_cpu_offload.py` covers automatic and forced streaming on nonzero devices.
+
+The SDXL stack (fp16 UNet, Canny ControlNet, both text encoders, the fp16-fix VAE)
+is ~8.8 GiB of weights, and the UNet and the ControlNet run together on every
+denoising step. Below the 12 GiB floor `_load_global` fuses the Lightning LoRA on
+the CPU and then calls `enable_sequential_cpu_offload`; the order matters, because
+the offload hooks stream their own copy of each weight and a later fuse would never
+reach it. Measured on an RTX 5060 Laptop (8 GiB, WSL2, 16 GiB system RAM) in
+October 2026, one 1672x941 global pass at strength 0.15:
+
+| Placement | Warm pass | Peak VRAM | Peak system RAM |
+|---|---|---|---|
+| resident `.to(device)` (previous behaviour) | did not finish in 10 min | 7.9 GiB used, spilling | -- |
+| `enable_model_cpu_offload` | 237 s | 9.2 GiB reserved on an 8 GiB card | 12.4 GiB |
+| group offload, block level, CUDA stream | 92 s, or hung | 12.4 GiB reserved | 15.6 GiB |
+| group offload, leaf level, CUDA stream | 7.8 s, or hung | 5.6 GiB | 15.7 GiB |
+| group offload, leaf level, `low_cpu_mem_usage` | 21.4 s | 5.6 GiB | 9.4 GiB |
+| `enable_sequential_cpu_offload` | 12.8 s | 3-4 GiB | 8.5 GiB |
+
+The resident and model-offload rows did not fail: under the Windows display driver
+CUDA falls back to shared system memory, so PyTorch reserved more than the card
+holds and the run crawled instead of raising. The streamed group-offload variants
+pin every weight in host memory up front; on a 16 GiB machine that left the run
+finishing or hanging depending on what else held RAM, which rules them out as a
+default. All completed placements produced bit-identical images.
+
+Fusing on the CPU instead of on CUDA rounds 788 of the UNet's 3256 fused tensors
+differently, by at most one fp16 ulp (1.2e-4). On the pass above that moved the
+output by a mean of 0.16 levels (PSNR 54.3 dB against a CUDA-fused, equally
+streamed run). Cards at or above the floor keep the CUDA fuse and therefore the
+previous pixels; only a streamed run, automatic or forced with `--cpu-offload`,
+sees the difference.
+
 Regression coverage:
 
 - [`test_cpu_offload.py`](../tests/test_cpu_offload.py)

src/remove_ai_watermarks/_internal/sdxl_zimage_pipeline.py

@@ -47,6 +47,12 @@ _LATENT_GRID = 8
 # The reference graph loads the Qwen LoRA at 0.8; carrying that number to a different
 # LoRA on a different architecture would be imitation, not parity.
 SDXL_STEPS = 4
+# Below this floor the global stack streams its weights from CPU layer by layer; at or
+# above it the stack stays resident. The fp16 UNet, Canny ControlNet, both text
+# encoders and the VAE are ~8.8 GiB of weights, and the UNet and the ControlNet run
+# together on every step, so model-level offload cannot get under an 8 GiB card.
+# Benchmark in docs/module-internals.md, "CPU offload".
+RESIDENT_SDXL_GLOBAL_MIN_VRAM_GIB = 12.0
 
 
 def sdxl_target_size(width: int, height: int) -> tuple[int, int]:
@@ -54,6 +60,17 @@ def sdxl_target_size(width: int, height: int) -> tuple[int, int]:
     return _target_size(width, height, _LATENT_GRID)
 
 
+def resolve_sdxl_global_residency(
+    requested: bool | None,
+    *,
+    total_memory_gib: float,
+) -> bool:
+    """Keep the SDXL global stack resident when explicitly requested or safely sized."""
+    if requested is not None:
+        return requested
+    return total_memory_gib >= RESIDENT_SDXL_GLOBAL_MIN_VRAM_GIB
+
+
 @dataclass
 class SdxlZImagePipeline(TwoStageZImagePipeline):
     """Lazy runtime for the SDXL global stage plus the inherited face stage."""
@@ -64,6 +81,12 @@ class SdxlZImagePipeline(TwoStageZImagePipeline):
         super().__post_init__()
         self._sdxl_pipe: Any = None
 
+    def _keep_global_models_resident(self) -> bool:
+        return resolve_sdxl_global_residency(
+            self.keep_global_models_on_device,
+            total_memory_gib=self._total_vram_gib(),
+        )
+
     def _load_global(self) -> Any:
         if self._sdxl_pipe is not None:
             return self._sdxl_pipe
@@ -89,11 +112,19 @@ class SdxlZImagePipeline(TwoStageZImagePipeline):
             variant="fp16",
             add_watermarker=False,
             **token,
-        ).to(self.device)
+        )
+        resident = self._keep_global_models_resident()
+        if resident:
+            pipe = pipe.to(self.device)
         pipe.load_lora_weights(hf_hub_download(SDXL_LIGHTNING_MODEL_ID, SDXL_LIGHTNING_PATTERN, **token))
         pipe.fuse_lora()
         # SDXL-Lightning is distilled against trailing timestep spacing.
         pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+        if not resident:
+            # After the fuse: the offload hooks keep their own CPU copy of each weight
+            # and stream that copy, so fusing later would never reach what runs.
+            self._progress("Streaming the SDXL stack from CPU (card below the residency floor)...")
+            pipe.enable_sequential_cpu_offload(device=self.device, gpu_id=torch.cuda.current_device())
         self._sdxl_pipe = pipe
         return pipe
 

src/remove_ai_watermarks/_internal/watermark_profiles.py

@@ -63,13 +63,13 @@ PROFILE_ADAPTIVE_POLISH = {
 
 # 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})
+# global stack is each profile's own: qwen-zimage streams it through
+# ``_qwen_vram_config`` and sdxl-zimage through Diffusers sequential offload.
+# chroma-zimage loads its stack with a plain ``.to(device)``, so on a card below the
+# face-stage residency floor the flag has nothing left to change there 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_ZIMAGE_PROFILE})
 
 SDXL_LIGHTNING_MODEL_ID = "ByteDance/SDXL-Lightning"
 SDXL_LIGHTNING_PATTERN = "sdxl_lightning_4step_lora.safetensors"

src/remove_ai_watermarks/_internal/watermark_remover.py

@@ -175,9 +175,9 @@ class WatermarkRemover:
             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.",
+            "qwen-zimage and sdxl-zimage stream their 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,
         )
 

src/remove_ai_watermarks/cli.py

@@ -308,8 +308,9 @@ _cpu_offload_option = click.option(
         "Offload model components to CPU between CUDA calls instead 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. It reaches "
-        "the global stack of qwen-zimage only: chroma-zimage and sdxl-zimage keep "
-        "theirs resident either way."
+        "the global stack of qwen-zimage and sdxl-zimage, which also stream theirs "
+        "automatically on a card too small to hold it; chroma-zimage keeps its "
+        "global stack resident either way."
     ),
 )
 

src/remove_ai_watermarks/invisible_engine.py

@@ -172,7 +172,8 @@ class InvisibleEngine:
             cpu_offload: Offload model components to CPU between CUDA calls instead
                 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.
+                every profile; reaches the global stack of qwen-zimage and
+                sdxl-zimage, not chroma-zimage.
                 CUDA only.
         """
 

tests/test_cpu_offload.py

@@ -1,20 +1,27 @@
 """Unit tests for how far --cpu-offload travels, per profile.
 
 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.
+reaches every profile; the global one is read by ``QwenZImagePipeline`` and
+``SdxlZImagePipeline``, so on ``chroma-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.
+needs no diffusion dependency, model download, or GPU. The SDXL loader placement is
+driven through stubbed Diffusers constructors and skips without the diffusion extra.
 """
 
 from __future__ import annotations
 
 import logging
+from typing import Any
+from unittest.mock import MagicMock
 
 import pytest
 
+from remove_ai_watermarks._internal.sdxl_zimage_pipeline import (
+    RESIDENT_SDXL_GLOBAL_MIN_VRAM_GIB,
+    resolve_sdxl_global_residency,
+)
 from remove_ai_watermarks._internal.watermark_profiles import (
     AUTO_PROFILE,
     CHROMA_ZIMAGE_PROFILE,
@@ -29,7 +36,7 @@ from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover
 OFFLOAD_SUPPORT_CASES = (
     (QWEN_ZIMAGE_PROFILE, True),
     (CHROMA_ZIMAGE_PROFILE, False),
-    (SDXL_ZIMAGE_PROFILE, False),
+    (SDXL_ZIMAGE_PROFILE, True),
     # Auto preloads its concrete Qwen fallback before an image vendor is known.
     (AUTO_PROFILE, True),
 )
@@ -69,18 +76,20 @@ class TestGlobalOffloadSupport:
 
 
 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:
+    def test_it_names_the_profile_and_the_flag(self, caplog: pytest.LogCaptureFixture) -> None:
         with caplog.at_level(logging.WARNING):
-            _remover(profile, cpu_offload=True)._warn_if_global_offload_unsupported()
+            _remover(CHROMA_ZIMAGE_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
+        assert CHROMA_ZIMAGE_PROFILE in message
 
-    def test_qwen_zimage_is_quiet(self, caplog: pytest.LogCaptureFixture) -> None:
+    @pytest.mark.parametrize("profile", [QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE])
+    def test_a_profile_that_streams_its_global_stack_is_quiet(
+        self, profile: str, caplog: pytest.LogCaptureFixture
+    ) -> None:
         with caplog.at_level(logging.WARNING):
-            _remover(QWEN_ZIMAGE_PROFILE, cpu_offload=True)._warn_if_global_offload_unsupported()
+            _remover(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:
@@ -132,3 +141,118 @@ class TestUnsupportedProfileWarns:
         with pytest.raises(ImportError):
             remover._load_qwen_zimage_pipeline()
         assert order == ["warned", "loaded"]
+
+
+class TestSdxlGlobalResidency:
+    @pytest.mark.parametrize(
+        ("requested", "total_memory_gib", "expected"),
+        [
+            # Unset: the card decides, and an 8 GiB card cannot hold the stack.
+            (None, 7.6, False),
+            (None, RESIDENT_SDXL_GLOBAL_MIN_VRAM_GIB, True),
+            (None, 23.5, True),
+            # An unreadable card reports 0.0 and must stream rather than spill.
+            (None, 0.0, False),
+            # Explicit requests win in both directions.
+            (False, 79.2, False),
+            (True, 7.6, True),
+        ],
+    )
+    def test_request_wins_and_the_card_decides_otherwise(
+        self, requested: bool | None, total_memory_gib: float, expected: bool
+    ) -> None:
+        assert resolve_sdxl_global_residency(requested, total_memory_gib=total_memory_gib) is expected
+
+
+class _RecordingSdxlPipe:
+    """Stands in for the Diffusers SDXL pipeline and records placement calls in order."""
+
+    def __init__(self, calls: list[Any]) -> None:
+        self.calls = calls
+        self.scheduler = MagicMock()
+
+    def to(self, device: str) -> _RecordingSdxlPipe:
+        self.calls.append(("to", device))
+        return self
+
+    def load_lora_weights(self, _path: str) -> None:
+        self.calls.append(("load_lora_weights",))
+
+    def fuse_lora(self) -> None:
+        self.calls.append(("fuse_lora",))
+
+    def enable_sequential_cpu_offload(self, *, device: str, gpu_id: int | None = None) -> None:
+        self.calls.append(("enable_sequential_cpu_offload", device, gpu_id))
+
+
+class TestSdxlGlobalPlacement:
+    @staticmethod
+    def _load(
+        monkeypatch: pytest.MonkeyPatch,
+        *,
+        requested: bool | None,
+        total_memory_gib: float,
+        current_device: int = 0,
+    ) -> list[Any]:
+        diffusers = pytest.importorskip("diffusers")
+        torch = pytest.importorskip("torch")
+        monkeypatch.setattr(torch.cuda, "current_device", lambda: current_device)
+        from remove_ai_watermarks._internal.sdxl_zimage_pipeline import SdxlZImagePipeline
+
+        calls: list[Any] = []
+        monkeypatch.setattr(diffusers.ControlNetModel, "from_pretrained", lambda *a, **k: MagicMock())
+        monkeypatch.setattr(diffusers.AutoencoderKL, "from_pretrained", lambda *a, **k: MagicMock())
+        monkeypatch.setattr(
+            diffusers.StableDiffusionXLControlNetImg2ImgPipeline,
+            "from_pretrained",
+            lambda *a, **k: _RecordingSdxlPipe(calls),
+        )
+        monkeypatch.setattr("huggingface_hub.hf_hub_download", lambda *a, **k: "lora.safetensors")
+        monkeypatch.setattr(diffusers.EulerDiscreteScheduler, "from_config", lambda *a, **k: MagicMock())
+
+        pipeline = SdxlZImagePipeline(device="cuda", torch_dtype=None, keep_global_models_on_device=requested)
+        monkeypatch.setattr(type(pipeline), "_require_cuda", lambda self: None)
+        monkeypatch.setattr(type(pipeline), "_total_vram_gib", lambda self: total_memory_gib)
+        pipeline._load_global()
+        return calls
+
+    def test_a_small_card_streams_the_fused_stack_instead_of_moving_it(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        calls = self._load(monkeypatch, requested=None, total_memory_gib=7.6)
+        # The offload hooks stream their own copy of each weight, so they must follow the fuse.
+        assert calls == [("load_lora_weights",), ("fuse_lora",), ("enable_sequential_cpu_offload", "cuda", 0)]
+
+    def test_a_large_card_keeps_the_stack_resident_as_before(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        calls = self._load(monkeypatch, requested=None, total_memory_gib=23.5)
+        assert calls == [("to", "cuda"), ("load_lora_weights",), ("fuse_lora",)]
+
+    def test_cpu_offload_streams_even_on_a_large_card(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        calls = self._load(monkeypatch, requested=False, total_memory_gib=79.2)
+        assert ("to", "cuda") not in calls
+        assert calls[-1] == ("enable_sequential_cpu_offload", "cuda", 0)
+
+    @pytest.mark.parametrize("current_device", [1, 2])
+    @pytest.mark.parametrize(("requested", "total_memory_gib"), [(None, 7.6), (False, 79.2)])
+    def test_streaming_uses_the_current_cuda_device(
+        self, monkeypatch: pytest.MonkeyPatch, current_device: int, requested: bool | None, total_memory_gib: float
+    ) -> None:
+        calls = self._load(
+            monkeypatch,
+            requested=requested,
+            total_memory_gib=total_memory_gib,
+            current_device=current_device,
+        )
+        assert ("to", "cuda") not in calls
+        assert calls[-1] == ("enable_sequential_cpu_offload", "cuda", current_device)
+
+    def test_cpu_offload_reaches_the_sdxl_stack(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        from remove_ai_watermarks._internal import sdxl_zimage_pipeline
+
+        captured: dict[str, object] = {}
+
+        class Recorder:
+            def __init__(self, **kwargs: object) -> None:
+                captured.update(kwargs)
+
+        monkeypatch.setattr(sdxl_zimage_pipeline, "SdxlZImagePipeline", Recorder)
+        _remover(SDXL_ZIMAGE_PROFILE, cpu_offload=True)._load_qwen_zimage_pipeline()
+        assert captured["keep_global_models_on_device"] is False

tests/test_platform.py

@@ -192,7 +192,7 @@ class TestNoReembeddedWatermark:
         # from_config would otherwise resolve the mock's config as a repo id.
         monkeypatch.setattr(diffusers.EulerDiscreteScheduler, "from_config", lambda *a, **k: MagicMock())
 
-        pipeline = SdxlZImagePipeline(device="cuda", torch_dtype=None)
+        pipeline = SdxlZImagePipeline(device="cuda", torch_dtype=None, keep_global_models_on_device=True)
         monkeypatch.setattr(type(pipeline), "_require_cuda", lambda self: None)
         pipeline._load_global()