From 36e37432893d234685509ad56515a6d4c75f4f76 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Sat, 1 Aug 2026 21:22:52 -0700 Subject: [PATCH] Keep the qwen-zimage global stack resident on a card that can hold it Port of the same change made on the v0.20.1 line, reapplied here because the package layout moved under _internal/ in the meantime. 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. Measured on an H100: a warm global pass went from 37.3s at 0.8 GiB resident to 2.2s at 28.7 GiB, with both stacks resident peaking at 48.0 GiB of 79.2. Co-Authored-By: Claude Opus 5 --- docs/module-internals.md | 23 +++++- .../_internal/qwen_zimage_pipeline.py | 59 ++++++++++++-- .../_internal/watermark_remover.py | 1 + tests/test_qwen_zimage_pipeline.py | 79 +++++++++++++++++++ 4 files changed, 155 insertions(+), 7 deletions(-) diff --git a/docs/module-internals.md b/docs/module-internals.md index 7caa042..af39a0f 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -518,11 +518,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) in +August 2026, 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 diff --git a/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py b/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py index ff9969e..5a7133f 100644 --- a/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py +++ b/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py @@ -55,6 +55,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 FACE_DENOISE_SCALE = 0.5 _CANNY_LOW = 13 @@ -79,6 +85,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"] @@ -497,6 +514,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: @@ -524,21 +542,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", diff --git a/src/remove_ai_watermarks/_internal/watermark_remover.py b/src/remove_ai_watermarks/_internal/watermark_remover.py index 891d904..e47763f 100644 --- a/src/remove_ai_watermarks/_internal/watermark_remover.py +++ b/src/remove_ai_watermarks/_internal/watermark_remover.py @@ -353,6 +353,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 diff --git a/tests/test_qwen_zimage_pipeline.py b/tests/test_qwen_zimage_pipeline.py index 2cf3a18..cdc120f 100644 --- a/tests/test_qwen_zimage_pipeline.py +++ b/tests/test_qwen_zimage_pipeline.py @@ -110,6 +110,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._internal.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 larger global stack pinned. + """ + from remove_ai_watermarks._internal import qwen_zimage_pipeline as pipeline_module + from remove_ai_watermarks._internal 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._internal.qwen_zimage_pipeline import ( QwenZImagePipeline,