Keep the qwen-zimage global stack resident on a card that can hold it

The mandatory Qwen stack was configured to offload to disk unconditionally.
DiffSynth implements that by dropping the weights to the meta device and
re-reading every parameter through its DiskMap on the next onload, and the
pipeline moves between text encoder, transformer and VAE on every pass, so
each generation paid a full model reload. That is the right trade on a
consumer card, where it is what makes a 20B model runnable at all, and pure
waste on a card that can simply hold the stack.

Residency is now resolved from total VRAM, mirroring how the optional
Z-Image face stack is already gated. Above the floor the config passes no
"disk" value anywhere, which is what actually disables the behavior:
DiffSynth latches disk_offload once from offload_dtype, so pointing every
device at CUDA while leaving the sentinel would keep both the meta-drop and
the re-read.

The floor is set equal to the face floor rather than lower because that is
the configuration measured with both stacks resident; a tighter gate is
plausible but unvalidated. cpu_offload now forces both stacks to stream, so
a caller asking for low VRAM no longer gets the larger stack pinned anyway.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

docs/module-internals.md

@@ -293,11 +293,32 @@ Regression coverage:
 
 CPU offload is enabled only when requested on CUDA. The standard Diffusers
 profiles call `enable_model_cpu_offload`. The `qwen-zimage` profile uses the
-same flag to force its face stack out of automatic device residency.
+same flag to force **both** its stacks out of automatic device residency.
+
+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`.
+
+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
+pure waste and the weights stay on the GPU. The difference is not marginal:
+DiffSynth offloads by dropping the weights to the meta device and re-reading every
+parameter through its `DiskMap` on the next onload, and the pipeline moves between
+text encoder, transformer and VAE on each pass. Measured on an H100 (80 GiB), a warm
+global pass took 37.3 s at 0.8 GiB resident with the streaming config, against 2.2 s
+at 28.7 GiB with the stack resident; both stacks resident peaked at 48.0 GiB. Faster
+storage cannot close that gap, because the cost is the reload itself rather than the
+read.
+
+The resident config deliberately passes no `"disk"` value anywhere. DiffSynth latches
+`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.
 
 Regression coverage:
 
 - [`test_cpu_offload.py`](../tests/test_cpu_offload.py)
+- [`test_qwen_zimage_pipeline.py`](../tests/test_qwen_zimage_pipeline.py)
 
 ### Qwen plus Z-Image
 

src/remove_ai_watermarks/noai/qwen_zimage_pipeline.py

@@ -65,6 +65,12 @@ GLOBAL_CFG = 1.0
 FACE_CFG = 1.0
 GLOBAL_CONTROLNET_SCALE = 1.0
 RESIDENT_FACE_MODEL_MIN_VRAM_GIB = 64.0
+# Below this floor the mandatory Qwen stack streams from disk, which is what makes a
+# 20B model runnable on a consumer card at all; at or above it, streaming is pure
+# waste. Set equal to the face floor rather than lower because that is the
+# configuration actually measured with both stacks resident; a tighter gate is
+# plausible but unvalidated. Benchmark in docs/module-internals.md, "CPU offload".
+RESIDENT_GLOBAL_MODEL_MIN_VRAM_GIB = 64.0
 # The reference face denoise assumes its ComfyUI detailer sampler, latent
 # noise-mask feather, and inpaint path. Applying that value unchanged to this
 # DiffSynth crop-regeneration port over-processes faces. Paired public-fixture
@@ -95,6 +101,17 @@ def resolve_face_model_residency(
     return total_memory_gib >= RESIDENT_FACE_MODEL_MIN_VRAM_GIB
 
 
+def resolve_global_model_residency(
+    requested: bool | None,
+    *,
+    total_memory_gib: float,
+) -> bool:
+    """Keep the mandatory Qwen stack resident when explicitly requested or safely sized."""
+    if requested is not None:
+        return requested
+    return total_memory_gib >= RESIDENT_GLOBAL_MODEL_MIN_VRAM_GIB
+
+
 def _pin_vram_managed_models(pipe: Any) -> None:
     """Move the managed Z-Image stack to CUDA once and make offload a no-op."""
     model_names = ["text_encoder", "dit", "vae_encoder", "vae_decoder"]
@@ -525,6 +542,7 @@ class QwenZImagePipeline:
     progress_callback: Callable[[str], None] | None = None
     controlnet_conditioning_scale: float = GLOBAL_CONTROLNET_SCALE
     keep_face_models_on_device: bool | None = None
+    keep_global_models_on_device: bool | None = None
     cache_prompt_embeddings: bool = True
 
     def __post_init__(self) -> None:
@@ -552,21 +570,50 @@ class QwenZImagePipeline:
             return max(1.0, torch.cuda.mem_get_info("cuda")[1] / (1024**3) - 0.5)
         return None
 
-    def _keep_face_models_resident(self) -> bool:
+    def _total_vram_gib(self) -> float:
+        """Card capacity for the residency gates, 0.0 when it cannot be read.
+
+        Separate from ``_vram_limit``, which answers a different question (the budget
+        handed to DiffSynth) and must stay ``None`` rather than 0.0 when unknown, so
+        an unreadable device means "no limit" there and "assume small" here.
+        """
         import torch
 
-        total_memory_gib = 0.0
         with contextlib.suppress(Exception):
-            total_memory_gib = torch.cuda.get_device_properties("cuda").total_memory / (1024**3)
+            return torch.cuda.get_device_properties("cuda").total_memory / (1024**3)
+        return 0.0
+
+    def _keep_face_models_resident(self) -> bool:
         return resolve_face_model_residency(
             self.keep_face_models_on_device,
-            total_memory_gib=total_memory_gib,
+            total_memory_gib=self._total_vram_gib(),
         )
 
-    @staticmethod
-    def _qwen_vram_config() -> dict[str, Any]:
+    def _keep_global_models_resident(self) -> bool:
+        return resolve_global_model_residency(
+            self.keep_global_models_on_device,
+            total_memory_gib=self._total_vram_gib(),
+        )
+
+    def _qwen_vram_config(self) -> dict[str, Any]:
         import torch
 
+        if self._keep_global_models_resident():
+            # Same fp8 storage and bf16 computation as the streaming config, but the
+            # weights never leave the GPU. Passing no "disk" anywhere is what makes
+            # this work: DiffSynth decides `disk_offload` once from `offload_dtype`,
+            # so a card large enough to hold the stack skips both the `to("meta")`
+            # drop and the DiskMap re-read entirely.
+            return {
+                "offload_dtype": torch.float8_e4m3fn,
+                "offload_device": "cuda",
+                "onload_dtype": torch.float8_e4m3fn,
+                "onload_device": "cuda",
+                "preparing_dtype": torch.float8_e4m3fn,
+                "preparing_device": "cuda",
+                "computation_dtype": torch.bfloat16,
+                "computation_device": "cuda",
+            }
         return {
             "offload_dtype": "disk",
             "offload_device": "disk",

src/remove_ai_watermarks/noai/watermark_remover.py

@@ -658,6 +658,7 @@ class WatermarkRemover:
                 progress_callback=self._progress_callback,
                 controlnet_conditioning_scale=self.controlnet_conditioning_scale,
                 keep_face_models_on_device=False if self.cpu_offload else None,
+                keep_global_models_on_device=False if self.cpu_offload else None,
             )
         return self._qwen_zimage_pipeline
 

tests/test_qwen_zimage_pipeline.py

@@ -101,6 +101,85 @@ def test_yunet_download_targets_verified_lfs_artifact():
     assert pytest.approx(0.5) == YUNET_SCORE_THRESHOLD
 
 
+def test_global_stack_is_resident_on_a_large_card_and_streams_on_a_small_one():
+    """The mandatory Qwen stack must not stream from disk on a card that can hold it.
+
+    DiffSynth offloads by dropping weights to the meta device and re-reading every
+    parameter through its DiskMap, so the streaming config costs a full model reload
+    on each stage transition. Benchmark in docs/module-internals.md, "CPU offload".
+    """
+    import torch
+
+    from remove_ai_watermarks.noai.qwen_zimage_pipeline import (
+        QwenZImagePipeline,
+        resolve_global_model_residency,
+    )
+
+    assert resolve_global_model_residency(None, total_memory_gib=79.2) is True
+    assert resolve_global_model_residency(None, total_memory_gib=39.5) is False
+    assert resolve_global_model_residency(False, total_memory_gib=79.2) is False
+    assert resolve_global_model_residency(True, total_memory_gib=39.5) is True
+
+    large = QwenZImagePipeline(
+        device="cuda",
+        torch_dtype=torch.bfloat16,
+        keep_global_models_on_device=True,
+    )._qwen_vram_config()
+    # No "disk" anywhere: DiffSynth latches disk_offload from offload_dtype once, so
+    # leaving it in would keep the meta-drop even with every device set to cuda.
+    assert "disk" not in large.values()
+    assert large["offload_device"] == "cuda"
+    assert large["onload_device"] == "cuda"
+    assert large["computation_dtype"] is torch.bfloat16
+
+    small = QwenZImagePipeline(
+        device="cuda",
+        torch_dtype=torch.bfloat16,
+        keep_global_models_on_device=False,
+    )._qwen_vram_config()
+    assert small["offload_dtype"] == "disk"
+    assert small["offload_device"] == "disk"
+    assert small["onload_device"] == "cpu"
+
+
+@pytest.mark.parametrize(
+    ("cpu_offload", "expected"),
+    [(True, False), (False, None)],
+)
+def test_cpu_offload_forces_both_stacks_to_stream(monkeypatch, cpu_offload, expected):
+    """``cpu_offload`` is the caller's escape hatch and must cover the global stack too.
+
+    Without the global flag it silenced only the face stack, so a caller asking for
+    low VRAM still got the 29 GiB global stack pinned.
+    """
+    from remove_ai_watermarks.noai import qwen_zimage_pipeline as pipeline_module
+    from remove_ai_watermarks.noai import watermark_remover as module
+
+    captured: dict[str, object] = {}
+
+    class Recorder:
+        def __init__(self, **kwargs):
+            captured.update(kwargs)
+
+    # `_load_qwen_zimage_pipeline` imports the class inside the function body, so the
+    # patch has to land on the defining module rather than on watermark_remover.
+    monkeypatch.setattr(pipeline_module, "QwenZImagePipeline", Recorder)
+
+    remover = module.WatermarkRemover.__new__(module.WatermarkRemover)
+    remover.device = "cuda"
+    remover.torch_dtype = None
+    remover.hf_token = None
+    remover._progress_callback = None
+    remover.controlnet_conditioning_scale = 1.0
+    remover.cpu_offload = cpu_offload
+    remover._qwen_zimage_pipeline = None
+
+    remover._load_qwen_zimage_pipeline()
+
+    assert captured["keep_global_models_on_device"] is expected
+    assert captured["keep_face_models_on_device"] is expected
+
+
 def test_resident_face_models_disable_vram_offload():
     from remove_ai_watermarks.noai.qwen_zimage_pipeline import (
         QwenZImagePipeline,