diff --git a/.claude/rules/development.md b/.claude/rules/development.md index 58f6d40..e76318d 100644 --- a/.claude/rules/development.md +++ b/.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. diff --git a/docs/module-internals.md b/docs/module-internals.md index 5bad63a..ff2acce 100644 --- a/docs/module-internals.md +++ b/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. diff --git a/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py b/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py index 54beb02..cc1ec4e 100644 --- a/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py +++ b/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( diff --git a/tests/test_qwen_zimage_pipeline.py b/tests/test_qwen_zimage_pipeline.py index 2da75e0..366a42d 100644 --- a/tests/test_qwen_zimage_pipeline.py +++ b/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,