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>

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(