Merge pull request #71 from wiltodelta/feat/static-prompt-embedding-cache

Serve the static prompt embeddings from a persistent cache

.claude/rules/development.md

@@ -24,7 +24,9 @@ Do not classify an entire module as untestable because its main path downloads a
 - target-size selection in `test_invisible_engine.py`;
 - unsharp and adaptive-polish helpers in `test_humanizer.py`;
 - mocked device fallback in `test_img2img_runner.py`;
-- tiling geometry and blending in `test_tiling.py`.
+- tiling geometry and blending in `test_tiling.py`;
+- prompt-embedding cache keying, storage round-trip, and the cross-pipeline reuse
+  that lets a stack load without its text encoder, in `test_qwen_zimage_pipeline.py`.
 
 Use availability checks only for paths that actually load large models.
 

docs/module-internals.md

@@ -564,6 +564,39 @@ orchestration, YuNet integration, SAM selection, masks, sizing helpers, and pixe
 compositing are implemented for this runtime. Changing a calibrated model input
 requires the same provider-oracle and identity evaluation as a model change.
 
+### Static prompt embeddings
+
+Both stages prompt with module constants, and at CFG 1.0 DiffSynth's
+`PipelineUnitRunner` reuses the positive embedding for the negative side instead of
+encoding it. So exactly one embedding per stage is ever computed, from text that
+cannot vary at runtime, which makes it cacheable across containers rather than only
+within one pipeline.
+
+`_cache_static_prompt_embeddings` therefore persists what the text encoder produced
+under `_model_cache_dir()/prompt-embeddings`, keyed by cache version, model id,
+pipeline output params, and the exact prompt string. Once that file exists,
+`_load_qwen` and `_load_zimage` drop the text-encoder `ModelConfig` from the model
+stack entirely and serve the stored tensors instead. Measured on an H100 volume in
+August 2026, that removes **15.45 GiB** (Qwen2.5-VL) and **7.49 GiB** (Z-Image) of a
+**87.6 GiB** per-request read, worth a median **11.76 s** and **4.10 s** of load time
+(paired within five containers). The output is **byte-identical** -- the stored
+tensors are the encoder's own -- so this needs no provider-oracle re-verification.
+
+Three properties are load-bearing:
+
+- **The key self-heals.** A model bump or a prompt edit changes the key, so the next
+  container recomputes rather than reading a stale embedding. `_PROMPT_CACHE_VERSION`
+  covers a change to the stored shape itself.
+- **The write is atomic.** A torn write must never be readable as a cache hit, so the
+  payload lands in a temp file and is renamed into place.
+- **A miss after the encoder was dropped raises.** `require_cache` records that the
+  stack was built without a text encoder on the strength of the file; falling back
+  would call a model that is not loaded, which surfaces as an opaque crash.
+
+`_model_cache_dir()` prefers `HF_HOME` for the same reason: on a scale-to-zero runner
+that is the only persistently mounted path, and anything below it is re-derived per
+request. The YuNet download follows the same root.
+
 Regression coverage:
 
 - [`test_qwen_zimage_pipeline.py`](../tests/test_qwen_zimage_pipeline.py)

src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py

@@ -63,6 +63,21 @@ RESIDENT_FACE_MODEL_MIN_VRAM_GIB = 64.0
 RESIDENT_GLOBAL_MODEL_MIN_VRAM_GIB = 64.0
 FACE_DENOISE_SCALE = 0.5
 
+# Both stages prompt with compile-time constants, and at CFG 1.0 DiffSynth reuses the
+# positive embedding for the negative side instead of encoding it, so exactly one
+# embedding per stage is ever needed. Persisting it lets the next container load the
+# stack without its text encoder at all: 15.45 GiB for Qwen and 7.49 GiB for Z-Image
+# of a measured 87.6 GiB read per request. The stored tensors are what the encoder
+# itself produced, so the output stays byte-identical. Bump the version when the
+# stored shape changes; the model id and prompt are already part of the key, so a
+# model or prompt change invalidates itself.
+_PROMPT_CACHE_VERSION = 1
+_PROMPT_CACHE_DIRNAME = "prompt-embeddings"
+# DiffSynth identifies a pipeline unit by what it produces, so these tuples are how
+# the prompt stage is located inside each pipeline and how its cache file is keyed.
+_QWEN_PROMPT_OUTPUTS = ("prompt_emb", "prompt_emb_mask")
+_ZIMAGE_PROMPT_OUTPUTS = ("prompt_embeds",)
+
 _CANNY_LOW = 13
 _CANNY_HIGH = 64
 
@@ -127,7 +142,56 @@ def _pin_vram_managed_models(pipe: Any) -> None:
     pipe.load_models_to_device(model_names)
 
 
-def _cached_prompt_process(original_process: Any) -> Any:
+def _cast_prompt_payload(payload: Any, device: Any, dtype: Any) -> Any:
+    """Move a stored embedding payload onto the runtime device.
+
+    Only floating tensors take the pipeline dtype. The attention masks travel in the
+    same payload and are integer, so casting them would corrupt the prompt.
+    """
+    import torch
+
+    if isinstance(payload, torch.Tensor):
+        if dtype is not None and payload.is_floating_point():
+            return payload.to(device=device, dtype=dtype)
+        return payload.to(device=device)
+    if isinstance(payload, dict):
+        return {key: _cast_prompt_payload(value, device, dtype) for key, value in payload.items()}
+    if isinstance(payload, (list, tuple)):
+        return type(payload)(_cast_prompt_payload(value, device, dtype) for value in payload)
+    return payload
+
+
+def _prompt_cache_path(model_id: str, output_params: tuple[str, ...], prompt: str) -> Path:
+    """Locate the stored embedding for one model and one exact prompt string."""
+    key = "\x1f".join((str(_PROMPT_CACHE_VERSION), model_id, *output_params, prompt))
+    digest = hashlib.sha256(key.encode("utf-8")).hexdigest()[:32]
+    return _model_cache_dir() / _PROMPT_CACHE_DIRNAME / f"{digest}.pt"
+
+
+def _store_prompt_payload(path: Path, payload: Any) -> None:
+    """Write the embedding atomically so a torn write can never read back as a hit."""
+    import torch
+
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.NamedTemporaryFile(dir=path.parent, suffix=".pt", delete=False) as handle:
+        torch.save(_cast_prompt_payload(payload, "cpu", None), handle)
+        temporary = Path(handle.name)
+    temporary.replace(path)
+
+
+def _load_prompt_payload(path: Path, device: Any, dtype: Any) -> Any:
+    import torch
+
+    return _cast_prompt_payload(torch.load(path, map_location="cpu", weights_only=True), device, dtype)
+
+
+def _cached_prompt_process(
+    original_process: Any,
+    *,
+    model_id: str | None,
+    output_params: tuple[str, ...],
+    require_cache: bool,
+) -> Any:
     cache: dict[str, dict[str, Any]] = {}
 
     def cached_process(
@@ -137,9 +201,30 @@ def _cached_prompt_process(original_process: Any) -> Any:
     ) -> dict[str, Any]:
         if edit_image is not None:
             return original_process(runtime_pipe, prompt, edit_image=edit_image)
-        if prompt not in cache:
-            cache[prompt] = original_process(runtime_pipe, prompt, edit_image=None)
-        return cache[prompt]
+        if prompt in cache:
+            return cache[prompt]
+        path = None if model_id is None else _prompt_cache_path(model_id, output_params, prompt)
+        if path is not None and path.exists():
+            try:
+                cache[prompt] = _load_prompt_payload(path, runtime_pipe.device, runtime_pipe.torch_dtype)
+            except Exception:
+                log.warning("Discarding unreadable prompt-embedding cache %s", path, exc_info=True)
+            else:
+                return cache[prompt]
+        if require_cache:
+            # The text encoder was left out of the model stack on the strength of
+            # this file, so there is nothing left to fall back to.
+            raise RuntimeError(
+                f"The cached prompt embedding {path} disappeared after the stack was loaded without its text encoder."
+            )
+        produced = original_process(runtime_pipe, prompt, edit_image=None)
+        if path is not None:
+            try:
+                _store_prompt_payload(path, produced)
+            except OSError:
+                log.warning("Could not persist the prompt embedding to %s", path, exc_info=True)
+        cache[prompt] = produced
+        return produced
 
     return cached_process
 
@@ -147,11 +232,25 @@ def _cached_prompt_process(original_process: Any) -> Any:
 def _cache_static_prompt_embeddings(
     pipe: Any,
     output_params: tuple[str, ...],
+    *,
+    model_id: str | None = None,
+    require_cache: bool = False,
 ) -> bool:
-    """Memoize a prompt unit when its embedding depends only on static text."""
+    """Memoize a prompt unit when its embedding depends only on static text.
+
+    With ``model_id`` the memo is also persisted, which is what lets the next
+    container skip the text encoder entirely; without it the memo lives only as long
+    as the pipeline. ``require_cache`` says the stack was already built without a
+    text encoder, so a miss must fail loudly rather than call a model that is absent.
+    """
     for unit in pipe.units:
         if tuple(getattr(unit, "output_params", ())) == output_params:
-            unit.process = _cached_prompt_process(unit.process)
+            unit.process = _cached_prompt_process(
+                unit.process,
+                model_id=model_id,
+                output_params=output_params,
+                require_cache=require_cache,
+            )
             return True
     return False
 
@@ -317,7 +416,13 @@ def _expanded_box(
 
 
 def _model_cache_dir() -> Path:
-    root = os.environ.get("XDG_CACHE_HOME")
+    """Where this library persists downloaded and derived model assets.
+
+    ``HF_HOME`` comes first because on a scale-to-zero runner it is the one path
+    mounted persistently; anything under a container-local cache is re-derived on
+    every request, which defeats both the YuNet download and the prompt cache.
+    """
+    root = os.environ.get("HF_HOME") or os.environ.get("XDG_CACHE_HOME")
     base = Path(root) if root else Path.home() / ".cache"
     return base / "remove-ai-watermarks"
 
@@ -567,6 +672,10 @@ class QwenZImagePipeline:
             total_memory_gib=self._total_vram_gib(),
         )
 
+    def _prompt_is_cached(self, model_id: str, output_params: tuple[str, ...], prompt: str) -> bool:
+        """Whether this stack can be built without its text encoder at all."""
+        return self.cache_prompt_embeddings and _prompt_cache_path(model_id, output_params, prompt).exists()
+
     def _qwen_vram_config(self) -> dict[str, Any]:
         import torch
 
@@ -629,17 +738,18 @@ class QwenZImagePipeline:
 
         self._progress("Loading Qwen-Image-2512, Lightning LoRA, and Canny ControlNet...")
         config = self._qwen_vram_config()
+        text_encoder_config = ModelConfig(
+            model_id=QWEN_IMAGE_2512_MODEL_ID,
+            origin_file_pattern="text_encoder/model*.safetensors",
+            **config,
+        )
         model_configs = [
             ModelConfig(
                 model_id=QWEN_IMAGE_2512_MODEL_ID,
                 origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors",
                 **config,
             ),
-            ModelConfig(
-                model_id=QWEN_IMAGE_2512_MODEL_ID,
-                origin_file_pattern="text_encoder/model*.safetensors",
-                **config,
-            ),
+            text_encoder_config,
             ModelConfig(
                 model_id=QWEN_IMAGE_2512_MODEL_ID,
                 origin_file_pattern="vae/diffusion_pytorch_model.safetensors",
@@ -651,6 +761,10 @@ class QwenZImagePipeline:
                 **config,
             ),
         ]
+        prompt_cached = self._prompt_is_cached(QWEN_IMAGE_2512_MODEL_ID, _QWEN_PROMPT_OUTPUTS, _GLOBAL_PROMPT)
+        if prompt_cached:
+            model_configs.remove(text_encoder_config)
+            log.info("Qwen prompt embedding is cached; loading the stack without its text encoder")
         pipe = QwenImagePipeline.from_pretrained(
             torch_dtype=self.torch_dtype,
             device=self.device,
@@ -675,7 +789,9 @@ class QwenZImagePipeline:
         if self.cache_prompt_embeddings:
             _cache_static_prompt_embeddings(
                 pipe,
-                ("prompt_emb", "prompt_emb_mask"),
+                _QWEN_PROMPT_OUTPUTS,
+                model_id=QWEN_IMAGE_2512_MODEL_ID,
+                require_cache=prompt_cached,
             )
         self._qwen_pipe = (pipe, ControlNetInput)
         return self._qwen_pipe
@@ -698,26 +814,32 @@ class QwenZImagePipeline:
         self._progress("Loading Z-Image Turbo face-detail model...")
         keep_on_device = self._keep_face_models_resident()
         config = self._zimage_vram_config()
+        text_encoder_config = ModelConfig(
+            model_id=ZIMAGE_TURBO_MODEL_ID,
+            origin_file_pattern="text_encoder/*.safetensors",
+            **config,
+        )
+        model_configs = [
+            ModelConfig(
+                model_id=ZIMAGE_TURBO_MODEL_ID,
+                origin_file_pattern="transformer/*.safetensors",
+                **config,
+            ),
+            text_encoder_config,
+            ModelConfig(
+                model_id=ZIMAGE_TURBO_MODEL_ID,
+                origin_file_pattern="vae/diffusion_pytorch_model.safetensors",
+                **config,
+            ),
+        ]
+        prompt_cached = self._prompt_is_cached(ZIMAGE_TURBO_MODEL_ID, _ZIMAGE_PROMPT_OUTPUTS, _FACE_PROMPT)
+        if prompt_cached:
+            model_configs.remove(text_encoder_config)
+            log.info("Z-Image prompt embedding is cached; loading the stack without its text encoder")
         pipe = ZImagePipeline.from_pretrained(
             torch_dtype=self.torch_dtype,
             device=self.device,
-            model_configs=[
-                ModelConfig(
-                    model_id=ZIMAGE_TURBO_MODEL_ID,
-                    origin_file_pattern="transformer/*.safetensors",
-                    **config,
-                ),
-                ModelConfig(
-                    model_id=ZIMAGE_TURBO_MODEL_ID,
-                    origin_file_pattern="text_encoder/*.safetensors",
-                    **config,
-                ),
-                ModelConfig(
-                    model_id=ZIMAGE_TURBO_MODEL_ID,
-                    origin_file_pattern="vae/diffusion_pytorch_model.safetensors",
-                    **config,
-                ),
-            ],
+            model_configs=model_configs,
             tokenizer_config=ModelConfig(
                 model_id=ZIMAGE_TURBO_MODEL_ID,
                 origin_file_pattern="tokenizer/",
@@ -727,7 +849,12 @@ class QwenZImagePipeline:
         if keep_on_device:
             _pin_vram_managed_models(pipe)
         if self.cache_prompt_embeddings:
-            _cache_static_prompt_embeddings(pipe, ("prompt_embeds",))
+            _cache_static_prompt_embeddings(
+                pipe,
+                _ZIMAGE_PROMPT_OUTPUTS,
+                model_id=ZIMAGE_TURBO_MODEL_ID,
+                require_cache=prompt_cached,
+            )
         self._zimage_pipe = pipe
         return pipe
 

tests/test_qwen_zimage_pipeline.py

@@ -271,6 +271,113 @@ def test_static_prompt_cache_reuses_embeddings_without_caching_image_edits():
     assert unit.calls == 4
 
 
+def test_prompt_cache_path_is_keyed_by_version_model_outputs_and_prompt(monkeypatch, tmp_path):
+    """A model, prompt or format change must not read a stale embedding."""
+    from remove_ai_watermarks._internal import qwen_zimage_pipeline as qz
+
+    monkeypatch.setenv("HF_HOME", str(tmp_path))
+    baseline = qz._prompt_cache_path("model/a", ("prompt_emb",), "text")
+
+    assert baseline.parent == tmp_path / "remove-ai-watermarks" / "prompt-embeddings"
+    assert baseline == qz._prompt_cache_path("model/a", ("prompt_emb",), "text")
+    assert baseline != qz._prompt_cache_path("model/b", ("prompt_emb",), "text")
+    assert baseline != qz._prompt_cache_path("model/a", ("prompt_embeds",), "text")
+    assert baseline != qz._prompt_cache_path("model/a", ("prompt_emb",), "other")
+
+    monkeypatch.setattr(qz, "_PROMPT_CACHE_VERSION", qz._PROMPT_CACHE_VERSION + 1)
+    assert baseline != qz._prompt_cache_path("model/a", ("prompt_emb",), "text")
+
+
+def test_model_cache_dir_prefers_the_persistent_hugging_face_root(monkeypatch, tmp_path):
+    """A scale-to-zero runner only mounts HF_HOME, so it must win over XDG."""
+    from remove_ai_watermarks._internal.qwen_zimage_pipeline import _model_cache_dir
+
+    monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg"))
+    monkeypatch.delenv("HF_HOME", raising=False)
+    assert _model_cache_dir() == tmp_path / "xdg" / "remove-ai-watermarks"
+
+    monkeypatch.setenv("HF_HOME", str(tmp_path / "hf"))
+    assert _model_cache_dir() == tmp_path / "hf" / "remove-ai-watermarks"
+
+
+def test_stored_prompt_embedding_round_trips_without_casting_the_mask(monkeypatch, tmp_path):
+    """The mask rides in the same payload and is integer; casting it corrupts the prompt."""
+    import torch
+
+    from remove_ai_watermarks._internal import qwen_zimage_pipeline as qz
+
+    monkeypatch.setenv("HF_HOME", str(tmp_path))
+    path = qz._prompt_cache_path("model/a", qz._QWEN_PROMPT_OUTPUTS, "text")
+    payload = {
+        "prompt_emb": torch.ones((1, 2, 3), dtype=torch.float32),
+        "prompt_emb_mask": torch.ones((1, 2), dtype=torch.int64),
+    }
+    qz._store_prompt_payload(path, payload)
+    restored = qz._load_prompt_payload(path, "cpu", torch.bfloat16)
+
+    assert restored["prompt_emb"].dtype == torch.bfloat16
+    assert restored["prompt_emb_mask"].dtype == torch.int64
+    assert torch.equal(restored["prompt_emb"].float(), payload["prompt_emb"])
+
+
+def test_persisted_prompt_cache_lets_a_second_pipeline_skip_the_text_encoder(monkeypatch, tmp_path):
+    """The whole point: container two must not call the encoder container one ran."""
+    import torch
+
+    from remove_ai_watermarks._internal import qwen_zimage_pipeline as qz
+
+    monkeypatch.setenv("HF_HOME", str(tmp_path))
+
+    class PromptUnit:
+        output_params = qz._ZIMAGE_PROMPT_OUTPUTS
+
+        def __init__(self):
+            self.calls = 0
+
+        def process(self, _pipe, prompt, edit_image=None):
+            self.calls += 1
+            return {"prompt_embeds": [torch.ones((2, 2), dtype=torch.float32)]}
+
+    def build():
+        unit = PromptUnit()
+        pipe = MagicMock(units=[unit], device="cpu", torch_dtype=torch.float32)
+        return unit, pipe
+
+    first_unit, first_pipe = build()
+    qz._cache_static_prompt_embeddings(first_pipe, qz._ZIMAGE_PROMPT_OUTPUTS, model_id="model/a", require_cache=False)
+    first_unit.process(first_pipe, qz._FACE_PROMPT)
+    assert first_unit.calls == 1
+
+    second_unit, second_pipe = build()
+    qz._cache_static_prompt_embeddings(second_pipe, qz._ZIMAGE_PROMPT_OUTPUTS, model_id="model/a", require_cache=True)
+    restored = second_unit.process(second_pipe, qz._FACE_PROMPT)
+
+    assert second_unit.calls == 0
+    assert torch.equal(restored["prompt_embeds"][0], torch.ones((2, 2)))
+
+
+def test_a_missing_cache_fails_loudly_once_the_text_encoder_was_left_out(monkeypatch, tmp_path):
+    """Silently calling an absent text encoder would surface as an opaque crash."""
+    import torch
+
+    from remove_ai_watermarks._internal import qwen_zimage_pipeline as qz
+
+    monkeypatch.setenv("HF_HOME", str(tmp_path))
+
+    class PromptUnit:
+        output_params = qz._QWEN_PROMPT_OUTPUTS
+
+        def process(self, _pipe, prompt, edit_image=None):
+            raise AssertionError("the text encoder is not loaded")
+
+    unit = PromptUnit()
+    pipe = MagicMock(units=[unit], device="cpu", torch_dtype=torch.float32)
+    qz._cache_static_prompt_embeddings(pipe, qz._QWEN_PROMPT_OUTPUTS, model_id="model/a", require_cache=True)
+
+    with pytest.raises(RuntimeError, match="disappeared"):
+        unit.process(pipe, "never cached")
+
+
 def test_sam_pixels_match_model_dtype_without_casting_boxes():
     import torch