Load the face stage in its own dtype, not the global stage's

The sdxl-zimage profile crashed on every image containing a face. The remover
gives it torch.float16, because SDXL ships fp16 weights and an fp16-safe VAE,
and that dtype reached the inherited _load_zimage while _zimage_vram_config
hardcodes bfloat16 for its offload, onload and computation dtypes. Z-Image was
therefore built bf16 and handed fp16 latents, dying in the VAE with "Input type
(c10::Half) and bias type (c10::BFloat16) should be the same". Zero-face inputs
never enter _run_faces, so the profile passed every timing run it was given, and
its tests avoid model downloads, so nothing exercised the loader.

Every face-stage loader now reads _face_stage_dtype(), the computation dtype of
the VRAM config it is paired with. SAM is routed through it too: it never
crashed, since it casts its own inputs and leaves through .float(), but it read
the same field and would have re-landed the bug for the next profile with a
different global dtype. That field was never the global dtype on this profile
anyway - _load_sdxl hardcodes fp16 for its own ControlNet, VAE and pipeline - so
its only readers were face-stage code.

This also fixes a second instance transitively: the persisted prompt-embedding
cache restores payloads at the DiffSynth pipe's dtype, which was fp16 into a
bf16 stack before this change.

For qwen-zimage the whole change is a strict no-op. The remover already hands it
bfloat16, the same value _face_stage_dtype() returns, so production is
untouched; verified on an H100 against the deployed pin.

The guard asserts the dtype the Z-Image and SAM loaders actually receive rather
than comparing the accessor to the config it derives from, which would restate
the implementation and pass for any consistently wrong value. Both assertions
were mutation-tested against the pre-fix line.

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

.claude/rules/development.md

@@ -26,7 +26,12 @@ Do not classify an entire module as untestable because its main path downloads a
 - mocked device fallback in `test_img2img_runner.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`.
+  that lets a stack load without its text encoder, in `test_qwen_zimage_pipeline.py`;
+- the face stack's dtype, in `test_qwen_zimage_pipeline.py`. A subclass that changes
+  the pipeline dtype for its own global model must not change the inherited face
+  stage's; `sdxl-zimage` shipped doing exactly that and crashed on every image with a
+  face. When one profile inherits another's stage, guard the invariants that stage
+  relies on, not just the code path.
 
 Use availability checks only for paths that actually load large models.
 

docs/module-internals.md

@@ -617,6 +617,38 @@ sigma range; Diffusers img2img truncates the step *count*
 (`init_timestep = int(steps * strength)`), so asking it for four steps at 0.15 executes
 **zero** and returns a bare VAE round-trip.
 
+**The face stage keeps its own dtype, and 0.23.0 shipped without that.** The remover
+gives this profile `torch.float16`, because SDXL ships fp16 weights and an fp16-safe
+VAE. That dtype reached the inherited `_load_zimage`, while `_zimage_vram_config()`
+hardcodes bfloat16 for its offload, onload and computation dtypes -- so the Z-Image
+modules were built bf16 and handed fp16 latents, and every image containing a face
+died in the VAE with `Input type (c10::Half) and bias type (c10::BFloat16) should be
+the same`. Every face-stage loader now reads `_face_stage_dtype()`, which returns the
+computation dtype of the VRAM config it is paired with, so the two cannot drift again.
+
+Two things hid this. Zero-face inputs never enter `_run_faces`, so the profile looked
+healthy on exactly the images used to time it; and the profile's tests deliberately
+avoid model downloads, so nothing exercised the loader. The lesson is narrower than
+"add a GPU test": inheriting a stage means inheriting its *invariants*, and this one
+was a dtype the subclass silently changed out from under it.
+
+Note what the seam is, because it decides where the fix belongs.
+`SdxlZImagePipeline._load_sdxl` hardcodes fp16 for its own ControlNet, VAE and
+pipeline, so `self.torch_dtype` was never actually the global stage's dtype on this
+profile -- its only remaining readers were face-stage code. SAM was the second one:
+it never crashed, because it casts its own inputs and leaves through `.float()`, but it
+was reading the same wrong field and would have re-landed the bug for the next profile
+with a different global dtype. It is routed through the same accessor, which for
+`qwen-zimage` is the bfloat16 it already used.
+
+The guard is `test_face_stage_loads_in_its_own_dtype_when_the_global_stage_differs`.
+It asserts the dtype the Z-Image and SAM loaders actually receive, not the accessor
+against the config it is derived from -- that comparison would restate the
+implementation and pass for any consistently wrong value. Both assertions were
+mutation-tested against the pre-fix line. For `qwen-zimage` the whole change is a
+strict no-op: the remover already handed it bfloat16, the same value
+`_face_stage_dtype()` returns.
+
 This profile is not deployed. Before it could be, it needs the other three Gemini
 originals, OpenAI re-verified at 0.15, a flat-graphic content class, and a low
 resolution case -- every verdict so far comes from one fixture and one seed.

src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py

@@ -706,6 +706,23 @@ class QwenZImagePipeline:
             "computation_device": "cuda",
         }
 
+    @classmethod
+    def _face_stage_dtype(cls) -> Any:
+        """The dtype every face-stage model loads and computes in.
+
+        Deliberately independent of ``self.torch_dtype``, which belongs to the global
+        stage. ``sdxl-zimage`` runs its global model in fp16, and inheriting that here
+        built the Z-Image modules bf16 (per the VRAM config below) while handing them
+        fp16 latents -- a Half/BFloat16 conv mismatch that crashed every face image
+        while zero-face inputs passed, so no unit test could see it.
+
+        Read by Z-Image and by SAM alike, so the whole stage moves together. SAM never
+        crashed, because it casts its own inputs and leaves through ``.float()``, but it
+        was reading the same wrong field and would re-land the bug for the next profile
+        that changes its global dtype.
+        """
+        return cls._zimage_vram_config()["computation_dtype"]
+
     @staticmethod
     def _zimage_vram_config() -> dict[str, Any]:
         import torch
@@ -837,7 +854,7 @@ class QwenZImagePipeline:
             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,
+            torch_dtype=self._face_stage_dtype(),
             device=self.device,
             model_configs=model_configs,
             tokenizer_config=ModelConfig(
@@ -872,7 +889,7 @@ class QwenZImagePipeline:
         processor = AutoProcessor.from_pretrained(SAM_MODEL_ID, **kwargs)
         model = AutoModelForMaskGeneration.from_pretrained(
             SAM_MODEL_ID,
-            torch_dtype=self.torch_dtype,
+            torch_dtype=self._face_stage_dtype(),
             **kwargs,
         ).to(self.device)
         model.eval()
@@ -901,7 +918,7 @@ class QwenZImagePipeline:
             )
             original_sizes = inputs["original_sizes"].clone()
             reshaped_sizes = inputs["reshaped_input_sizes"].clone()
-            inputs = _prepare_sam_inputs(inputs, self.device, self.torch_dtype)
+            inputs = _prepare_sam_inputs(inputs, self.device, self._face_stage_dtype())
             with torch.inference_mode():
                 outputs = model(**inputs, multimask_output=True)
             processed = processor.post_process_masks(

tests/test_qwen_zimage_pipeline.py

@@ -189,6 +189,50 @@ def test_cpu_offload_forces_both_stacks_to_stream(monkeypatch, cpu_offload, expe
     assert captured["keep_face_models_on_device"] is expected
 
 
+def test_face_stage_loads_in_its_own_dtype_when_the_global_stage_differs(monkeypatch, tmp_path):
+    """A subclass that changes the pipeline dtype must not change the face stage's.
+
+    ``sdxl-zimage`` is constructed fp16 for its global model. That dtype used to reach
+    the inherited ``_load_zimage``, which builds its modules bf16 from
+    ``_zimage_vram_config``, so Z-Image got fp16 latents into bf16 convolutions and
+    every image containing a face died in the VAE. Zero-face inputs never enter the
+    face stage, so the profile looked healthy right up to the first portrait.
+
+    Asserts the dtype the loaders actually RECEIVE. Comparing the accessor against the
+    config it is derived from would restate the implementation and pass for any
+    consistently-wrong value.
+    """
+    import torch
+    import transformers
+    from diffsynth.pipelines import z_image
+
+    from remove_ai_watermarks._internal.sdxl_zimage_pipeline import SdxlZImagePipeline
+
+    monkeypatch.setenv("HF_HOME", str(tmp_path))
+    captured: dict[str, object] = {}
+
+    def fake_zimage(**kwargs):
+        captured["zimage"] = kwargs["torch_dtype"]
+        return MagicMock(units=[])
+
+    def fake_sam(_model_id, **kwargs):
+        captured["sam"] = kwargs["torch_dtype"]
+        return MagicMock()
+
+    monkeypatch.setattr(z_image.ZImagePipeline, "from_pretrained", staticmethod(fake_zimage))
+    monkeypatch.setattr(transformers.AutoProcessor, "from_pretrained", staticmethod(lambda *a, **k: MagicMock()))
+    monkeypatch.setattr(transformers.AutoModelForMaskGeneration, "from_pretrained", staticmethod(fake_sam))
+
+    pipeline = SdxlZImagePipeline(device="cuda", torch_dtype=torch.float16)
+    pipeline._load_zimage()
+    pipeline._load_sam()
+
+    assert pipeline.torch_dtype == torch.float16, "the global stage keeps its own dtype"
+    # Z-Image is the one that crashed; SAM never did, but it read the same wrong field.
+    assert captured["zimage"] == torch.bfloat16
+    assert captured["sam"] == torch.bfloat16
+
+
 def test_resident_face_models_disable_vram_offload():
     from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
         QwenZImagePipeline,