From a64d81f446e5802aa2d4963678454901570820a1 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Sat, 1 Aug 2026 21:06:13 -0700 Subject: [PATCH] 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 --- docs/module-internals.md | 23 +++++- .../noai/qwen_zimage_pipeline.py | 59 ++++++++++++-- .../noai/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 3e9d0e4..b978686 100644 --- a/docs/module-internals.md +++ b/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 diff --git a/src/remove_ai_watermarks/noai/qwen_zimage_pipeline.py b/src/remove_ai_watermarks/noai/qwen_zimage_pipeline.py index 25df2f0..bbc3af5 100644 --- a/src/remove_ai_watermarks/noai/qwen_zimage_pipeline.py +++ b/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", diff --git a/src/remove_ai_watermarks/noai/watermark_remover.py b/src/remove_ai_watermarks/noai/watermark_remover.py index c1ae671..9ab3710 100644 --- a/src/remove_ai_watermarks/noai/watermark_remover.py +++ b/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 diff --git a/tests/test_qwen_zimage_pipeline.py b/tests/test_qwen_zimage_pipeline.py index cb3d520..ddc4be4 100644 --- a/tests/test_qwen_zimage_pipeline.py +++ b/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,