diff --git a/docs/module-internals.md b/docs/module-internals.md index 467538b..5bad63a 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -589,6 +589,56 @@ 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. +### SDXL plus Z-Image + +[`_internal/sdxl_zimage_pipeline.py`](../src/remove_ai_watermarks/_internal/sdxl_zimage_pipeline.py) +runs the same two-stage recipe on an SDXL global pass. `SdxlZImagePipeline` subclasses +`QwenZImagePipeline` and overrides only `_run_global` and `preload`, so the face stage +is inherited rather than copied and cannot drift between the profiles; a test asserts +the shared methods are the same objects. + +Four things are architecture-bound and swap with the model: the ControlNet +(`xinsir/controlnet-canny-sdxl-1.0`), the four-step distillation LoRA +(`ByteDance/SDXL-Lightning` at its documented strength 1.0, not the reference graph's +0.8, which belongs to a different LoRA), the sampler (Euler with trailing spacing, no +AuraFlow shift), and the latent grid (8 px against Qwen's 16). + +**Strength is architecture-bound too, and that is the easy mistake.** An SDXL global +pass leaves SynthID at the strength Qwen needs: verified through the Gemini app on a +native 2816x1536 original, 0.154 is FOUND while 0.20, 0.25 and 0.30 are clean. So this +profile takes a vendor policy (`SDXL_ZIMAGE_OPENAI_STRENGTH` 0.15, +`SDXL_ZIMAGE_GEMINI_STRENGTH` 0.25, unknown following Gemini) rather than +`resolution_adaptive_denoise`. Flat values are what was measured; no size dependence +has been established for this stage, so none is asserted. + +`requested_steps` exists because the two runtimes truncate differently. DiffSynth sets +`sigma_start = denoising_strength` and runs every requested step across the shortened +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. + +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. + +### Measured provider boundaries for qwen-zimage + +Both ends of the shipped curve now have oracle verdicts, and the shipped curve clears +everything it has been tested at: + +| oracle | fixture size | detected at | clean from | +|---|---|---|---| +| openai.com/verify | 1.57 MP | 0.06 | 0.08 | +| Gemini app | 4.33 MP | 0.08 | 0.10 | +| Gemini app | 0.57 MP | -- | 0.0896 (the curve's own value) | +| Gemini app | 1.40 MP | -- | 0.1066 (the curve's own value) | + +Read the last two rows before concluding the curve's low end is under-driven. Against +the 4.33 MP Gemini boundary the sub-1 MP rungs of 0.084-0.094 look short, but at those +sizes the curve's own values verify clean, which is what a resolution-scaled +requirement would predict. There is no measured size at which the shipped curve fails, +so it is left alone. + ### Static prompt embeddings Both stages prompt with module constants, and at CFG 1.0 DiffSynth's diff --git a/src/remove_ai_watermarks/_internal/sdxl_zimage_pipeline.py b/src/remove_ai_watermarks/_internal/sdxl_zimage_pipeline.py new file mode 100644 index 0000000..f72ac50 --- /dev/null +++ b/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 diff --git a/src/remove_ai_watermarks/_internal/watermark_profiles.py b/src/remove_ai_watermarks/_internal/watermark_profiles.py index 22e99f8..850f19c 100644 --- a/src/remove_ai_watermarks/_internal/watermark_profiles.py +++ b/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) diff --git a/src/remove_ai_watermarks/_internal/watermark_remover.py b/src/remove_ai_watermarks/_internal/watermark_remover.py index e47763f..9dc6067 100644 --- a/src/remove_ai_watermarks/_internal/watermark_remover.py +++ b/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) diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index 317bcac..b7d34fa 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/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, diff --git a/tests/test_qwen_zimage_pipeline.py b/tests/test_qwen_zimage_pipeline.py index a33fa30..2da75e0 100644 --- a/tests/test_qwen_zimage_pipeline.py +++ b/tests/test_qwen_zimage_pipeline.py @@ -728,3 +728,74 @@ def test_invisible_engine_uses_qwen_zimage_step_default(tmp_image_path, tmp_path ) assert engine._remover.remove_watermark.call_args.kwargs["num_inference_steps"] == 4 + + +def test_sdxl_zimage_strength_is_vendor_adaptive_and_leaves_other_profiles_alone(): + """An SDXL global pass needs more strength than Qwen, so it gets its own policy.""" + from remove_ai_watermarks._internal.watermark_profiles import ( + SDXL_ZIMAGE_GEMINI_STRENGTH, + SDXL_ZIMAGE_OPENAI_STRENGTH, + resolve_strength, + ) + + assert resolve_strength(None, "openai", "sdxl-zimage") == pytest.approx(SDXL_ZIMAGE_OPENAI_STRENGTH) + assert resolve_strength(None, "google", "sdxl-zimage") == pytest.approx(SDXL_ZIMAGE_GEMINI_STRENGTH) + # Unknown provenance takes the stricter of the two. + assert resolve_strength(None, None, "sdxl-zimage") == pytest.approx(SDXL_ZIMAGE_GEMINI_STRENGTH) + # An explicit value still wins, and the older profiles are untouched. + assert resolve_strength(0.4, "google", "sdxl-zimage") == pytest.approx(0.4) + assert resolve_strength(None, "openai", "controlnet") == pytest.approx(0.10) + assert resolve_strength(None, "google", "controlnet") == pytest.approx(0.15) + + +def test_sdxl_zimage_shares_the_four_step_seed_and_step_contract(): + from remove_ai_watermarks._internal.watermark_profiles import ( + normalize_profile, + resolve_seed, + resolve_steps, + ) + + assert normalize_profile("sdxl_zimage") == "sdxl-zimage" + assert resolve_steps(None, "sdxl-zimage") == 4 + assert resolve_seed(None, "sdxl-zimage") == 0 + assert resolve_steps(None, "controlnet") == 50 + assert resolve_seed(None, "controlnet") is None + + +def test_sdxl_requested_steps_compensate_for_the_diffusers_truncation(): + """Diffusers truncates the step COUNT where DiffSynth truncates the sigma range. + + Asking Diffusers for four steps at strength 0.15 runs int(4 * 0.15) = 0 and + returns a bare VAE round-trip, so the request has to be scaled up instead. + """ + from remove_ai_watermarks._internal.sdxl_zimage_pipeline import requested_steps + + for strength in (0.15, 0.25, 0.0896): + steps = requested_steps(4, strength) + assert int(steps * strength) >= 4 + # Naively asking for four would have under-spent every time, and at the + # strengths this profile actually uses it would have run nothing at all. + assert int(4 * strength) < 4 + assert int(4 * 0.15) == 0 + + +def test_sdxl_zimage_floors_to_its_own_latent_grid(): + """SDXL aligns to 8 pixels where Qwen aligns to 16.""" + from remove_ai_watermarks._internal.qwen_zimage_pipeline import _target_size + from remove_ai_watermarks._internal.sdxl_zimage_pipeline import sdxl_target_size + + assert sdxl_target_size(1122, 1402) == (1120, 1400) + assert _target_size(1122, 1402) == (1120, 1392) + assert sdxl_target_size(3, 3) == (8, 8) + + +def test_sdxl_zimage_inherits_the_face_stage_rather_than_copying_it(): + """The face stage must not be able to diverge between the two profiles.""" + from remove_ai_watermarks._internal.qwen_zimage_pipeline import QwenZImagePipeline + from remove_ai_watermarks._internal.sdxl_zimage_pipeline import SdxlZImagePipeline + + assert issubclass(SdxlZImagePipeline, QwenZImagePipeline) + for shared in ("_run_faces", "_sam_masks", "_load_zimage", "_load_sam", "run"): + assert getattr(SdxlZImagePipeline, shared) is getattr(QwenZImagePipeline, shared) + # Only the global stage and what it needs may differ. + assert SdxlZImagePipeline._run_global is not QwenZImagePipeline._run_global