Add the sdxl-zimage profile: the same recipe on an SDXL global stage

SdxlZImagePipeline subclasses QwenZImagePipeline and overrides only _run_global
and preload, so the face stage is inherited rather than copied and cannot drift
between the two profiles. A test asserts the shared methods are the same objects.

Four things are architecture-bound and swap with the model: the ControlNet, the
four-step distillation LoRA (SDXL-Lightning at its documented 1.0, not the
reference graph's 0.8, which belongs to a different LoRA), the sampler (Euler
trailing, no AuraFlow shift), and the latent grid at 8 px against Qwen's 16.

Strength is architecture-bound too, which is the easy mistake and cost two wrong
conclusions before it was caught. An SDXL global pass leaves SynthID at the
strength Qwen needs: through the Gemini app on a native 2816x1536 original, 0.154
is FOUND while 0.20, 0.25 and 0.30 are clean. The profile therefore takes a flat
vendor policy - OpenAI 0.15, Gemini 0.25, unknown following Gemini - rather than
resolution_adaptive_denoise, because flat values are what was measured and no
size dependence has been established for this stage.

requested_steps exists because the runtimes truncate differently: DiffSynth sets
sigma_start = denoising_strength and runs every requested step, while Diffusers
img2img truncates the step count, so four steps at 0.15 executes zero and returns
a bare VAE round-trip.

Also records both measured provider boundaries for the shipped qwen-zimage curve
- OpenAI detected at 0.06 and clean from 0.08, Gemini detected at 0.08 and clean
from 0.10 - together with the two low-resolution Gemini verdicts that explain why
the curve's sub-1 MP rungs are not under-driven despite looking short against a
boundary measured at 4.33 MP. The curve is left unchanged; nothing measured fails.

The profile is not deployed and not production-ready: every verdict so far comes
from one fixture and one seed.

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

src/remove_ai_watermarks/_internal/sdxl_zimage_pipeline.py

@@ -0,0 +1,154 @@
+"""The qwen-zimage recipe with an SDXL global stage.
+
+Only the global regeneration model changes. The face stage is inherited verbatim
+from :class:`QwenZImagePipeline` -- same YuNet detection, same SAM masks, same
+Z-Image Turbo repair of the original crops, same feathered compositing -- so a
+change there cannot silently diverge between the two profiles.
+
+Three pieces cannot be shared, because they are bound to the architecture: the
+ControlNet, the four-step distillation LoRA, and the sampler. Strength is bound to
+it too, which is the part that is easy to miss: an SDXL global pass leaves SynthID
+at the strength Qwen needs. See ``watermark_profiles.SDXL_ZIMAGE_OPENAI_STRENGTH``.
+"""
+
+# Diffusers and torch expose mostly untyped tensor APIs. Keep the relaxation local
+# to this optional ML boundary.
+# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportAttributeAccessIssue=false, reportPrivateUsage=false, reportPrivateImportUsage=false
+from __future__ import annotations
+
+import logging
+import math
+from dataclasses import dataclass
+from typing import Any
+
+from PIL import Image
+
+from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
+    GLOBAL_STEPS,
+    QwenZImagePipeline,
+    build_canny_control_image,
+)
+from remove_ai_watermarks._internal.watermark_profiles import (
+    CONTROLNET_CANNY_MODEL,
+    DEFAULT_MODEL_ID,
+    SDXL_LIGHTNING_MODEL_ID,
+    SDXL_LIGHTNING_PATTERN,
+)
+
+log = logging.getLogger(__name__)
+
+SDXL_VAE_MODEL_ID = "madebyollin/sdxl-vae-fp16-fix"
+# SDXL aligns to an 8-pixel latent grid, against Qwen's 16.
+_LATENT_GRID = 8
+
+
+def sdxl_target_size(width: int, height: int) -> tuple[int, int]:
+    """Floor dimensions to SDXL's latent grid without changing aspect."""
+    return max(_LATENT_GRID, (width // _LATENT_GRID) * _LATENT_GRID), max(
+        _LATENT_GRID, (height // _LATENT_GRID) * _LATENT_GRID
+    )
+
+
+def requested_steps(effective_steps: int, strength: float) -> int:
+    """Translate "spend N denoising steps" into what Diffusers has to be asked for.
+
+    The two runtimes truncate differently and it is easy to port this wrong.
+    DiffSynth sets ``sigma_start = denoising_strength`` and then runs *every*
+    requested step across the shortened sigma range. Diffusers img2img instead
+    truncates the step *count* (``init_timestep = int(steps * strength)``), so
+    asking it for four steps at strength 0.15 runs **zero** and returns nothing but
+    a VAE round-trip. Ask for enough that ``effective_steps`` actually execute.
+    """
+    return max(1, math.ceil(effective_steps / max(float(strength), 1e-6)))
+
+
+@dataclass
+class SdxlZImagePipeline(QwenZImagePipeline):
+    """Lazy runtime for the SDXL global stage plus the inherited face stage."""
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self._sdxl_pipe: Any = None
+
+    def _load_sdxl(self) -> Any:
+        if self._sdxl_pipe is not None:
+            return self._sdxl_pipe
+        self._require_cuda()
+        import torch
+        from diffusers import (
+            AutoencoderKL,
+            ControlNetModel,
+            EulerDiscreteScheduler,
+            StableDiffusionXLControlNetImg2ImgPipeline,
+        )
+        from huggingface_hub import hf_hub_download
+
+        self._progress("Loading SDXL, Lightning LoRA, and Canny ControlNet...")
+        token = {"token": self.hf_token} if self.hf_token else {}
+        controlnet = ControlNetModel.from_pretrained(CONTROLNET_CANNY_MODEL, torch_dtype=torch.float16, **token)
+        vae = AutoencoderKL.from_pretrained(SDXL_VAE_MODEL_ID, torch_dtype=torch.float16, **token)
+        pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
+            DEFAULT_MODEL_ID,
+            controlnet=controlnet,
+            vae=vae,
+            torch_dtype=torch.float16,
+            variant="fp16",
+            add_watermarker=False,
+            **token,
+        ).to(self.device)
+        # SDXL's own four-step distillation, at the strength its authors document.
+        # The reference graph loads the Qwen LoRA at 0.8; carrying that number to a
+        # different LoRA on a different architecture would be imitation, not parity.
+        pipe.load_lora_weights(hf_hub_download(SDXL_LIGHTNING_MODEL_ID, SDXL_LIGHTNING_PATTERN, **token))
+        pipe.fuse_lora()
+        # SDXL-Lightning is distilled against trailing timestep spacing.
+        pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+        self._sdxl_pipe = pipe
+        return pipe
+
+    def preload(self, *, global_only: bool = False) -> None:
+        """Eagerly load the mandatory stage and, by default, the face stack."""
+        from remove_ai_watermarks._internal.qwen_zimage_pipeline import _yunet_model_path
+
+        self._load_sdxl()
+        _yunet_model_path()
+        if not global_only:
+            self._load_zimage()
+            self._load_sam()
+
+    def _run_global(self, image: Image.Image, strength: float, seed: int | None) -> Image.Image:
+        import torch
+
+        pipe = self._load_sdxl()
+        target = sdxl_target_size(image.width, image.height)
+        prepared = image if image.size == target else image.resize(target, Image.Resampling.LANCZOS)
+        control = build_canny_control_image(prepared)
+        steps = requested_steps(GLOBAL_STEPS, strength)
+        self._progress(f"Running SDXL Canny pass: strength={strength:.4f}, steps={GLOBAL_STEPS} of {steps}...")
+        generator = torch.Generator(device=self.device).manual_seed(seed) if seed is not None else None
+        result = pipe(
+            prompt=self._global_prompt(),
+            negative_prompt=self._global_negative(),
+            image=prepared,
+            control_image=control,
+            controlnet_conditioning_scale=float(self.controlnet_conditioning_scale),
+            strength=float(strength),
+            num_inference_steps=steps,
+            guidance_scale=1.0,
+            generator=generator,
+        ).images[0]
+        if result.size != image.size:
+            result = result.resize(image.size, Image.Resampling.LANCZOS)
+        return result.convert("RGB")
+
+    @staticmethod
+    def _global_prompt() -> str:
+        from remove_ai_watermarks._internal.qwen_zimage_pipeline import _GLOBAL_PROMPT
+
+        return _GLOBAL_PROMPT
+
+    @staticmethod
+    def _global_negative() -> str:
+        from remove_ai_watermarks._internal.qwen_zimage_pipeline import _GLOBAL_NEGATIVE
+
+        return _GLOBAL_NEGATIVE

src/remove_ai_watermarks/_internal/watermark_profiles.py

@@ -15,6 +15,10 @@ CONTROLNET_CANNY_MODEL = "xinsir/controlnet-canny-sdxl-1.0"
 
 SDXL_PROFILE = "sdxl"
 QWEN_ZIMAGE_PROFILE = "qwen-zimage"
+SDXL_ZIMAGE_PROFILE = "sdxl-zimage"
+
+SDXL_LIGHTNING_MODEL_ID = "ByteDance/SDXL-Lightning"
+SDXL_LIGHTNING_PATTERN = "sdxl_lightning_4step_lora.safetensors"
 
 OPENAI_STRENGTH = 0.10
 GEMINI_STRENGTH = 0.15
@@ -25,6 +29,22 @@ QWEN_OPENAI_STRENGTH = 0.10
 QWEN_GEMINI_STRENGTH = 0.25
 QWEN_UNKNOWN_STRENGTH = QWEN_GEMINI_STRENGTH
 
+# sdxl-zimage runs the qwen-zimage recipe on an SDXL global stage, and strength is
+# architecture-bound: at Qwen's 0.154 an SDXL global pass leaves SynthID on a native
+# 2816x1536 Gemini original, while 0.20, 0.25 and 0.30 all read clean in the Gemini
+# app. 0.25 keeps a rung of margin over that boundary, which the historical SDXL
+# certification argues for -- it recorded 0.20 as DETECTED against Gemini on an
+# older SDXL pipeline. OpenAI is the easier oracle: the profile already cleared
+# openai.com/verify at 0.1102, so 0.15 sits above what was verified rather than on
+# it. Unknown follows Gemini, the stricter of the two.
+#
+# Unlike qwen-zimage this is a flat vendor policy rather than a resolution curve,
+# because flat values are what was measured. Every verdict above comes from a fixed
+# strength at one size; no size dependence has been established for this stage.
+SDXL_ZIMAGE_OPENAI_STRENGTH = 0.15
+SDXL_ZIMAGE_GEMINI_STRENGTH = 0.25
+SDXL_ZIMAGE_UNKNOWN_STRENGTH = SDXL_ZIMAGE_GEMINI_STRENGTH
+
 
 @dataclass(frozen=True)
 class _StrengthPolicy:
@@ -43,7 +63,16 @@ _QWEN_POLICY = _StrengthPolicy(
     unknown=QWEN_UNKNOWN_STRENGTH,
     by_vendor={"openai": QWEN_OPENAI_STRENGTH, "google": QWEN_GEMINI_STRENGTH},
 )
-_ALIASES = {"default": SDXL_PROFILE, "qwen_zimage": QWEN_ZIMAGE_PROFILE}
+_SDXL_ZIMAGE_POLICY = _StrengthPolicy(
+    unknown=SDXL_ZIMAGE_UNKNOWN_STRENGTH,
+    by_vendor={"openai": SDXL_ZIMAGE_OPENAI_STRENGTH, "google": SDXL_ZIMAGE_GEMINI_STRENGTH},
+)
+_ALIASES = {
+    "default": SDXL_PROFILE,
+    "qwen_zimage": QWEN_ZIMAGE_PROFILE,
+    "sdxl_zimage": SDXL_ZIMAGE_PROFILE,
+}
+_FOUR_STEP_PROFILES = frozenset({QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE})
 
 
 def normalize_profile(profile: str) -> str:
@@ -56,14 +85,14 @@ def resolve_steps(num_inference_steps: int | None, pipeline: str) -> int:
     """Return an explicit step count or the selected profile's default."""
     if num_inference_steps is not None:
         return num_inference_steps
-    return 4 if normalize_profile(pipeline) == QWEN_ZIMAGE_PROFILE else 50
+    return 4 if normalize_profile(pipeline) in _FOUR_STEP_PROFILES else 50
 
 
 def resolve_seed(seed: int | None, pipeline: str) -> int | None:
-    """Keep the fixed Qwen plus Z-Image profile reproducible by default."""
+    """Keep the fixed four-step Z-Image profiles reproducible by default."""
     if seed is not None:
         return seed
-    return 0 if normalize_profile(pipeline) == QWEN_ZIMAGE_PROFILE else None
+    return 0 if normalize_profile(pipeline) in _FOUR_STEP_PROFILES else None
 
 
 def strength_default_help() -> str:
@@ -79,7 +108,13 @@ def resolve_strength(strength: float | None, vendor: str | None = None, pipeline
     """Resolve a user override or the calibrated policy for a profile and vendor."""
     if strength is not None:
         return strength
-    policy = _QWEN_POLICY if pipeline is not None and normalize_profile(pipeline) == "qwen" else _STANDARD_POLICY
+    profile = normalize_profile(pipeline) if pipeline is not None else ""
+    if profile == "qwen":
+        policy = _QWEN_POLICY
+    elif profile == SDXL_ZIMAGE_PROFILE:
+        policy = _SDXL_ZIMAGE_POLICY
+    else:
+        policy = _STANDARD_POLICY
     return policy.choose(vendor)
 
 

src/remove_ai_watermarks/_internal/watermark_remover.py

@@ -17,6 +17,7 @@ from remove_ai_watermarks._internal.watermark_profiles import (
     DEFAULT_STRENGTH,
     QWEN_MODEL_ID,
     QWEN_ZIMAGE_PROFILE,
+    SDXL_ZIMAGE_PROFILE,
     normalize_profile,
     resolve_seed,
     resolve_steps,
@@ -30,6 +31,14 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+# Both two-stage profiles share the face stage, the four-step schedule, CFG 1.0, the
+# fixed model stack and the native-resolution contract; only the global model differs.
+_ZIMAGE_PROFILES = frozenset({QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE})
+_ZIMAGE_STACKS = {
+    QWEN_ZIMAGE_PROFILE: "Qwen-Image-2512 and Z-Image",
+    SDXL_ZIMAGE_PROFILE: "SDXL and Z-Image",
+}
+
 try:
     import torch
 
@@ -195,8 +204,10 @@ class WatermarkRemover:
     ) -> None:
         requested_model = model_id or self.DEFAULT_MODEL_ID
         self.model_profile = normalize_profile(pipeline)
-        if self.model_profile == QWEN_ZIMAGE_PROFILE and model_id not in {None, self.DEFAULT_MODEL_ID}:
-            raise ValueError("The qwen-zimage profile uses a fixed Qwen-Image-2512 and Z-Image model stack.")
+        if self.model_profile in _ZIMAGE_PROFILES and model_id not in {None, self.DEFAULT_MODEL_ID}:
+            raise ValueError(
+                f"The {self.model_profile} profile uses a fixed {_ZIMAGE_STACKS[self.model_profile]} model stack."
+            )
         self.model_id = (
             "Qwen/Qwen-Image-2512 + Tongyi-MAI/Z-Image-Turbo"
             if self.model_profile == QWEN_ZIMAGE_PROFILE
@@ -212,6 +223,10 @@ class WatermarkRemover:
             self.torch_dtype = torch_dtype
         elif self.device in {"cpu", "mps"}:
             self.torch_dtype = torch.float32  # type: ignore[union-attr]
+        elif self.model_profile == SDXL_ZIMAGE_PROFILE:
+            # SDXL ships fp16 weights and an fp16-safe VAE; bf16 would give up the
+            # variant without buying anything on this architecture.
+            self.torch_dtype = torch.float16  # type: ignore[union-attr]
         elif self.model_profile in {"qwen", QWEN_ZIMAGE_PROFILE}:
             self.torch_dtype = torch.bfloat16  # type: ignore[union-attr]
         else:
@@ -233,7 +248,7 @@ class WatermarkRemover:
 
     def preload(self, *, global_only: bool = False) -> None:
         """Materialize the selected model stack before the first request."""
-        if self.model_profile == QWEN_ZIMAGE_PROFILE:
+        if self.model_profile in _ZIMAGE_PROFILES:
             self._load_qwen_zimage_pipeline().preload(global_only=global_only)
         elif self.model_profile == "qwen":
             self._load_qwen_pipeline()
@@ -344,9 +359,16 @@ class WatermarkRemover:
 
     def _load_qwen_zimage_pipeline(self) -> Any:
         if self._qwen_zimage_pipeline is None:
-            from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline
+            if getattr(self, "model_profile", QWEN_ZIMAGE_PROFILE) == SDXL_ZIMAGE_PROFILE:
+                from remove_ai_watermarks._internal.sdxl_zimage_pipeline import (
+                    SdxlZImagePipeline as _Pipeline,
+                )
+            else:
+                from remove_ai_watermarks._internal.qwen_zimage_pipeline import (
+                    QwenZImagePipeline as _Pipeline,
+                )
 
-            self._qwen_zimage_pipeline = QwenZImagePipeline(
+            self._qwen_zimage_pipeline = _Pipeline(
                 device=self.device,
                 torch_dtype=self.torch_dtype,
                 hf_token=self.hf_token,
@@ -471,7 +493,7 @@ class WatermarkRemover:
         tile_size: int,
         tile_overlap: int,
     ) -> Image.Image:
-        if self.model_profile == QWEN_ZIMAGE_PROFILE:
+        if self.model_profile in _ZIMAGE_PROFILES:
             return self._run_qwen_zimage(
                 image,
                 strength,
@@ -543,14 +565,12 @@ class WatermarkRemover:
 
         resolved_seed = resolve_seed(seed, self.model_profile)
         steps = resolve_steps(num_inference_steps, self.model_profile)
-        guidance = (
-            1.0 if guidance_scale is None and self.model_profile == QWEN_ZIMAGE_PROFILE else guidance_scale or 7.5
-        )
-        if self.model_profile == QWEN_ZIMAGE_PROFILE:
+        guidance = 1.0 if guidance_scale is None and self.model_profile in _ZIMAGE_PROFILES else guidance_scale or 7.5
+        if self.model_profile in _ZIMAGE_PROFILES:
             if steps != 4:
-                raise ValueError("The qwen-zimage profile requires 4 steps.")
+                raise ValueError(f"The {self.model_profile} profile requires 4 steps.")
             if guidance != 1.0:
-                raise ValueError("The qwen-zimage profile requires CFG 1.0.")
+                raise ValueError(f"The {self.model_profile} profile requires CFG 1.0.")
         else:
             steps = viable_steps(steps, resolved_strength)
 

src/remove_ai_watermarks/invisible_engine.py

@@ -107,7 +107,9 @@ class InvisibleEngine:
                 flat-graphic content), or "qwen" (Qwen-Image 20B img2img, best text/
                 structure preservation but CUDA/cloud-class), or "qwen-zimage"
                 (Qwen-Image-2512 Lightning + Canny, then SAM-masked Z-Image face
-                repair; CUDA-only). "default" aliases "sdxl".
+                repair; CUDA-only), or "sdxl-zimage" (the same recipe and the same face
+                stage on an SDXL global pass, vendor-adaptive strength because an SDXL
+                global stage needs more of it; CUDA-only). "default" aliases "sdxl".
             hf_token: HuggingFace API token.
             progress_callback: Optional callback for progress messages.
             controlnet_conditioning_scale: ControlNet structure-preservation
@@ -238,7 +240,7 @@ class InvisibleEngine:
 
         if num_inference_steps is None:
             profile = getattr(self._remover, "model_profile", None)
-            num_inference_steps = 4 if profile == "qwen-zimage" else 100
+            num_inference_steps = 4 if profile in {"qwen-zimage", "sdxl-zimage"} else 100
         profile = getattr(self._remover, "model_profile", "controlnet")
         seed = resolve_seed(seed, profile)
 
@@ -266,7 +268,7 @@ class InvisibleEngine:
         # Keep an explicit max cap available for callers, but do not apply the SDXL
         # 1024px minimum-resolution floor to this profile.
         effective_min_resolution = (
-            0 if getattr(self._remover, "model_profile", None) == "qwen-zimage" else min_resolution
+            0 if getattr(self._remover, "model_profile", None) in {"qwen-zimage", "sdxl-zimage"} else min_resolution
         )
         target = _target_size(
             image.width,
@@ -392,7 +394,7 @@ class InvisibleEngine:
         """Remove invisible watermarks from all images in a directory."""
         if steps is None:
             profile = getattr(self._remover, "model_profile", None)
-            steps = 4 if profile == "qwen-zimage" else 50
+            steps = 4 if profile in {"qwen-zimage", "sdxl-zimage"} else 50
         return self._remover.remove_watermark_batch(
             input_dir=input_dir,
             output_dir=output_dir,