diff --git a/.claude/rules/development.md b/.claude/rules/development.md index e76318d..280b34e 100644 --- a/.claude/rules/development.md +++ b/.claude/rules/development.md @@ -23,7 +23,6 @@ Do not classify an entire module as untestable because its main path downloads a - target-size selection in `test_invisible_engine.py`; - unsharp and adaptive-polish helpers in `test_humanizer.py`; -- 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`; diff --git a/docs/cli.md b/docs/cli.md index bee3c92..5e0630a 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -363,10 +363,15 @@ remove-ai-watermarks invisible image.png -o clean.png --force | Pipeline | When to use it | | --- | --- | -| `controlnet` | Default compatibility profile with structural conditioning | -| `sdxl` | Lighter plain SDXL regeneration | -| `qwen` | Large CUDA oriented Qwen Image profile | -| `qwen-zimage` | CUDA only high fidelity profile with a separate face stage | +| `qwen-zimage` | Default. Qwen-Image-2512 global pass plus a SAM-masked Z-Image face stage | +| `sdxl-zimage` | The same recipe and face stage on an SDXL global pass, at a higher denoise | + +**Both are CUDA-only.** There is no CPU or MPS profile for invisible-watermark +removal. The former `controlnet`, `sdxl`, `qwen` and `default` profiles were removed +rather than kept as a CPU path: none of them matched this recipe's face preservation, +so offering them implied a quality the library no longer delivers. Passing a retired +name is rejected at parse time rather than remapped. Visible-mark removal and every +identify path still run anywhere. Example: diff --git a/docs/python-api.md b/docs/python-api.md index 050877d..7d94872 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -389,7 +389,7 @@ from pathlib import Path from remove_ai_watermarks.invisible_engine import InvisibleEngine engine = InvisibleEngine( - pipeline="controlnet", + pipeline="qwen-zimage", # the default; the only other value is "sdxl-zimage" device=None, cpu_offload=False, ) @@ -407,15 +407,16 @@ For limited CUDA memory: ```python engine = InvisibleEngine( - pipeline="controlnet", + pipeline="qwen-zimage", cpu_offload=True, ) ``` -For the CUDA only high fidelity profile: +Both profiles are CUDA-only, so `device=None` resolving to CPU or MPS cannot run +invisible-watermark removal at all. For the SDXL global stage instead of Qwen: ```python -engine = InvisibleEngine(pipeline="qwen-zimage") +engine = InvisibleEngine(pipeline="sdxl-zimage") ``` The `qwen-zimage` extra must be installed for that profile. diff --git a/docs/synthid-robust-identity-research.md b/docs/synthid-robust-identity-research.md index 3877a63..5295c5a 100644 --- a/docs/synthid-robust-identity-research.md +++ b/docs/synthid-robust-identity-research.md @@ -246,7 +246,9 @@ from the test set + this doc). - If no -> the assumption is wrong; PhotoMaker would re-introduce the watermark. Stop and reconsider. 2. **PhotoMaker-V1 prototype** in the existing `controlnet` pipeline: - - Mirror the `_load_controlnet_pipeline` path: add a PhotoMaker variant that + - Mirror the `_load_controlnet_pipeline` path (removed in 0.24.0 with the + controlnet profile; the equivalent seam is now `SdxlZImagePipeline._load_sdxl`): + add a PhotoMaker variant that loads SDXL + canny ControlNet + PhotoMaker adapter on the same engine. - Extract the OpenCLIP face embedding from the watermarked face crops (use OpenCV YuNet, already bundled for `auto`, to find the face boxes). diff --git a/docs/synthid.md b/docs/synthid.md index d09faa7..66b4eae 100644 --- a/docs/synthid.md +++ b/docs/synthid.md @@ -514,10 +514,17 @@ study (section 2.2) gives empirical floors: - **Google native 2816**: 0.15 clears (n=2, deployed controlnet worker, 2026-06-14) -- the same rung as capped 1536, so no resolution penalty was observed. -The default is **vendor-adaptive** (`watermark_profiles.resolve_strength` + -`vendor_for_strength`): the tool reads the C2PA issuer on the original input and picks +> **Superseded in 0.24.0.** The `sdxl`, `controlnet`, `qwen` and `default` profiles +> were removed, and `OPENAI_STRENGTH` / `GEMINI_STRENGTH` / `UNKNOWN_STRENGTH` went +> with them. Everything from here to the end of this section is a record of what was +> measured on those profiles, kept because the oracle verdicts are still the evidence +> base. For the strength policy in force now see `module-internals.md`: `qwen-zimage` +> uses `resolution_adaptive_denoise`, `sdxl-zimage` a flat vendor ladder. + +The default was **vendor-adaptive** (`watermark_profiles.resolve_strength` + +`vendor_for_strength`): the tool read the C2PA issuer on the original input and picked `OPENAI_STRENGTH` 0.10 / `GEMINI_STRENGTH` 0.15 / `UNKNOWN_STRENGTH` 0.15 **(LOWERED -2026-06-14 from the 2026-06-04 cert floors 0.20/0.30/0.30)**. **The SAME ladder applies +2026-06-14 from the 2026-06-04 cert floors 0.20/0.30/0.30)**. **The SAME ladder applied to both pipelines** (`sdxl` and `controlnet`). The 2026-06-14 re-test on the deployed Modal controlnet worker (v0.10.0) cleared SynthID on the oracle at OpenAI 0.10 (2 photoreal) and Google 0.15 (2 NATIVE 2816x1536, contradicting the "native >= 0.30" guess diff --git a/packaging/conda/recipe.yaml b/packaging/conda/recipe.yaml index ae2b2a0..4945e88 100644 --- a/packaging/conda/recipe.yaml +++ b/packaging/conda/recipe.yaml @@ -1,7 +1,7 @@ schema_version: 1 context: - version: "0.23.0" + version: "0.24.0" python_min: "3.10" package: diff --git a/pyproject.toml b/pyproject.toml index b6a3aac..72ffde1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "remove-ai-watermarks" -version = "0.23.0" +version = "0.24.0" description = "AI watermark remover for visible, invisible, and provenance marks in images and video" readme = "README.md" requires-python = ">=3.10.1" diff --git a/src/remove_ai_watermarks/__init__.py b/src/remove_ai_watermarks/__init__.py index bcf6622..b6d0118 100644 --- a/src/remove_ai_watermarks/__init__.py +++ b/src/remove_ai_watermarks/__init__.py @@ -32,7 +32,7 @@ _os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error") _warnings.filterwarnings("ignore", message=r".*ImageProcessorFast.*") -__version__ = "0.23.0" +__version__ = "0.24.0" __all__ = [ "__version__", diff --git a/src/remove_ai_watermarks/_internal/img2img_runner.py b/src/remove_ai_watermarks/_internal/img2img_runner.py deleted file mode 100644 index 2fcd37c..0000000 --- a/src/remove_ai_watermarks/_internal/img2img_runner.py +++ /dev/null @@ -1,136 +0,0 @@ -"""Execute Diffusers img2img calls and recover from an MPS runtime failure.""" - -from __future__ import annotations - -import contextlib -import logging -from typing import TYPE_CHECKING, Any - -from remove_ai_watermarks._internal.progress import is_mps_error, make_pipeline_progress - -if TYPE_CHECKING: - from collections.abc import Callable - - from PIL import Image - -logger = logging.getLogger(__name__) - - -def _pipeline_arguments( - image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - step_callback: Any, - overrides: dict[str, Any] | None, -) -> dict[str, Any]: - arguments: dict[str, Any] = { - "prompt": "", - "image": image, - "strength": strength, - "num_inference_steps": num_inference_steps, - "guidance_scale": guidance_scale, - "generator": generator, - } - arguments.update(overrides or {}) - if step_callback is not None: - arguments.update(callback=step_callback, callback_steps=1) - return arguments - - -def _invoke(pipeline: Any, arguments: dict[str, Any]) -> Image.Image: - response = pipeline(**arguments) - return response.images[0] - - -def run_img2img( - pipeline: Any, - image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - device: str, - set_progress: Callable[[str], None], - extra_kwargs: dict[str, Any] | None = None, -) -> Image.Image: - """Run one img2img request and report denoising progress when supported.""" - callback, started, finished, launch_monitor = make_pipeline_progress( - max(1, int(num_inference_steps * strength)), device, set_progress - ) - launch_monitor() - arguments = _pipeline_arguments( - image, strength, num_inference_steps, guidance_scale, generator, callback, extra_kwargs - ) - try: - try: - return _invoke(pipeline, arguments) - except TypeError as error: - if "callback" not in str(error): - raise - started.set() - arguments.pop("callback", None) - arguments.pop("callback_steps", None) - return _invoke(pipeline, arguments) - finally: - started.set() - finished.set() - - -def run_img2img_with_mps_fallback( - load_pipeline: Callable[[], Any], - image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - device: str, - set_progress: Callable[[str], None], - *, - reload_on_cpu: Callable[[], Any], - extra_kwargs: dict[str, Any] | None = None, -) -> tuple[Image.Image, str]: - """Retry an MPS-specific failure once with a freshly loaded CPU pipeline.""" - try: - output = run_img2img( - load_pipeline(), - image, - strength, - num_inference_steps, - guidance_scale, - generator, - device, - set_progress, - extra_kwargs, - ) - return output, device - except RuntimeError as error: - if device != "mps" or not is_mps_error(error): - raise - logger.warning("MPS execution failed (%s); retrying on CPU", error) - set_progress("MPS execution failed; retrying on CPU...") - try_empty_device_cache("mps") - output = run_img2img( - reload_on_cpu(), - image, - strength, - num_inference_steps, - guidance_scale, - None, - "cpu", - set_progress, - extra_kwargs, - ) - return output, "cpu" - - -def try_empty_device_cache(device: str) -> None: - """Ask Torch to release cached accelerator memory when the backend supports it.""" - with contextlib.suppress(Exception): - import torch - - backend = getattr(torch, device, None) - empty_cache = getattr(backend, "empty_cache", None) - if callable(empty_cache): - empty_cache() diff --git a/src/remove_ai_watermarks/_internal/progress.py b/src/remove_ai_watermarks/_internal/progress.py deleted file mode 100644 index 2fb8f05..0000000 --- a/src/remove_ai_watermarks/_internal/progress.py +++ /dev/null @@ -1,212 +0,0 @@ -"""Progress reporting utilities for long-running optional model operations.""" - -from __future__ import annotations - -import contextlib -import io -import os -import sys -import threading -import time -import warnings -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from collections.abc import Callable - -_BAR_WIDTH = 28 -_SPINNER = ("|", "/", "-", "\\") - - -def _truncate(text: str, max_len: int = 72) -> str: - if len(text) <= max_len: - return text - return f"{text[: max(0, max_len - 3)]}..." - - -def _build_bar(step: int) -> str: - position = step % (2 * _BAR_WIDTH - 2) - if position >= _BAR_WIDTH: - position = 2 * _BAR_WIDTH - 2 - position - cells = ["-"] * _BAR_WIDTH - cells[position] = "=" - return "".join(cells) - - -@dataclass -class _TaskResult: - value: Any = None - error: BaseException | None = None - complete: threading.Event = field(default_factory=threading.Event) - - -def run_with_progress(task: Callable[[], Any], progress_state: dict[str, str] | None = None) -> Any: - """Run ``task`` on a worker thread and render a compact terminal heartbeat.""" - outcome = _TaskResult() - - def invoke() -> None: - try: - outcome.value = task() - except BaseException as error: # re-raised on the caller thread - outcome.error = error - finally: - outcome.complete.set() - - worker = threading.Thread(target=invoke, name="raiw-progress-task", daemon=True) - worker.start() - started_at = time.monotonic() - frame = 0 - terminal = sys.__stderr__ - while not outcome.complete.wait(0.1): - message = _truncate((progress_state or {}).get("message", "Processing...")) - elapsed = int(time.monotonic() - started_at) - if terminal is not None: - terminal.write( - f"\r\033[2K {_SPINNER[frame % len(_SPINNER)]} [{_build_bar(frame)}] {elapsed:>3}s {message}" - ) - terminal.flush() - frame += 1 - - worker.join() - elapsed = int(time.monotonic() - started_at) - message = _truncate((progress_state or {}).get("message", "Processing...")) - if terminal is not None: - terminal.write(f"\r\033[2K Completed in {elapsed}s {message}\n") - terminal.flush() - if outcome.error is not None: - raise outcome.error - return outcome.value - - -def _silence_diffusers() -> None: - from diffusers.utils import logging as diffusers_logging - - diffusers_logging.set_verbosity_error() - disable = getattr(diffusers_logging, "disable_progress_bar", None) - if callable(disable): - disable() - - -def _configure_quiet_libraries() -> None: - os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") - operations = ( - lambda: __import__("transformers").logging.set_verbosity_error(), - _silence_diffusers, - lambda: __import__("huggingface_hub").logging.set_verbosity_error(), - ) - for operation in operations: - with contextlib.suppress(Exception): - operation() - - -def silence_library_output( - run_func: Callable[[], Any], - set_progress: Callable[[str], None] | None = None, -) -> Callable[[], Any]: - """Wrap a model call so third-party progress bars do not corrupt our CLI UI.""" - - def quiet_call() -> Any: - if set_progress is not None: - set_progress("Preparing model runtime...") - _configure_quiet_libraries() - with ( - warnings.catch_warnings(), - contextlib.redirect_stdout(io.StringIO()), - contextlib.redirect_stderr(io.StringIO()), - ): - warnings.simplefilter("ignore") - if set_progress is not None: - set_progress("Running watermark regeneration...") - return run_func() - - return quiet_call - - -@dataclass -class _PipelineMonitor: - total_steps: int - device: str - update: Callable[[str], None] - bar_len: int - label: str - pre_phases: list[tuple[int, str]] - post_phases: list[tuple[int, str]] - first_step: threading.Event = field(default_factory=threading.Event) - done: threading.Event = field(default_factory=threading.Event) - started_at: float = field(default_factory=time.monotonic) - last_step_at: float = field(default_factory=time.monotonic) - - def callback(self, step: int, _timestep: int, _latents: Any) -> None: - self.first_step.set() - now = time.monotonic() - self.last_step_at = now - current = min(self.total_steps, step + 1) - filled = round(self.bar_len * current / self.total_steps) - elapsed = now - self.started_at - eta = elapsed * max(0, self.total_steps - current) / max(1, current) - bar = "#" * filled + "." * (self.bar_len - filled) - self.update( - f"{self.label} [{bar}] {current}/{self.total_steps}, " - f"{elapsed:.0f}s elapsed, ~{eta:.0f}s left, {self.device}" - ) - - def _phase_message(self, phases: list[tuple[int, str]], elapsed: float) -> str: - message = phases[0][1] - for threshold, candidate in phases: - if elapsed < threshold: - break - message = candidate - return message - - def monitor(self) -> None: - while not self.first_step.wait(0.4): - elapsed = time.monotonic() - self.started_at - self.update(self._phase_message(self.pre_phases, elapsed)) - decode_started: float | None = None - while not self.done.wait(0.4): - if time.monotonic() - self.last_step_at < 1.5: - decode_started = None - continue - decode_started = decode_started or time.monotonic() - self.update(self._phase_message(self.post_phases, time.monotonic() - decode_started)) - - def start(self) -> threading.Thread: - self.started_at = self.last_step_at = time.monotonic() - self.first_step.clear() - self.done.clear() - thread = threading.Thread(target=self.monitor, name="raiw-pipeline-progress", daemon=True) - thread.start() - return thread - - -def make_pipeline_progress( - effective_steps: int, - device: str, - set_progress: Callable[[str], None], - *, - bar_len: int = 20, - label: str = "Denoising", - pre_phases: list[tuple[int, str]] | None = None, - post_phases: list[tuple[int, str]] | None = None, -) -> tuple[Callable[..., None], threading.Event, threading.Event, Callable[[], threading.Thread]]: - """Build a callback and monitor for the legacy Diffusers callback interface.""" - - def qualify(entries: list[tuple[int, str]]) -> list[tuple[int, str]]: - return [(second, f"{text} on {device}") for second, text in entries] - - monitor = _PipelineMonitor( - total_steps=max(1, effective_steps), - device=device, - update=set_progress, - bar_len=bar_len, - label=label, - pre_phases=pre_phases or qualify([(0, "Encoding image"), (8, "Preparing denoiser"), (20, "Starting sampler")]), - post_phases=post_phases or qualify([(0, "Decoding image"), (10, "Finalizing pixels"), (45, "Still decoding")]), - ) - return monitor.callback, monitor.first_step, monitor.done, monitor.start - - -def is_mps_error(error: Exception) -> bool: - """Return whether an error message identifies Apple's MPS backend.""" - return "mps" in str(error).casefold() diff --git a/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py b/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py index cc1ec4e..90fd41e 100644 --- a/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py +++ b/src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py @@ -1039,7 +1039,7 @@ class QwenZImagePipeline: ) -> Image.Image: """Execute global regeneration and masked face repair.""" self._require_cuda() - seed = resolve_seed(seed, "qwen-zimage") + seed = resolve_seed(seed) global_strength = ( resolution_adaptive_denoise(image.width, image.height) if strength is None else float(strength) ) diff --git a/src/remove_ai_watermarks/_internal/watermark_profiles.py b/src/remove_ai_watermarks/_internal/watermark_profiles.py index 850f19c..fa4ee13 100644 --- a/src/remove_ai_watermarks/_internal/watermark_profiles.py +++ b/src/remove_ai_watermarks/_internal/watermark_profiles.py @@ -1,33 +1,39 @@ -"""Project-owned configuration for invisible-watermark regeneration profiles.""" +"""Project-owned configuration for invisible-watermark regeneration profiles. + +Two profiles remain, and both are CUDA-only: ``qwen-zimage`` (the default) and +``sdxl-zimage``. The older ``controlnet``, ``sdxl``, ``qwen`` and ``default`` profiles +were removed rather than kept as a CPU path, because none of them matched the two-stage +recipe's face preservation and keeping them implied a quality this library no longer +offers. Removing invisible watermarks therefore needs a CUDA device; the visible-mark +registry and every identify path still run anywhere. +""" from __future__ import annotations -import math from dataclasses import dataclass from typing import TYPE_CHECKING, Literal if TYPE_CHECKING: from pathlib import Path +# SDXL base is no longer a profile of its own, but it is still the global stage of +# sdxl-zimage, so the checkpoint id stays. DEFAULT_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0" -QWEN_MODEL_ID = "Qwen/Qwen-Image" CONTROLNET_CANNY_MODEL = "xinsir/controlnet-canny-sdxl-1.0" -SDXL_PROFILE = "sdxl" QWEN_ZIMAGE_PROFILE = "qwen-zimage" SDXL_ZIMAGE_PROFILE = "sdxl-zimage" +DEFAULT_PROFILE = QWEN_ZIMAGE_PROFILE +PROFILE_CHOICES = (QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE) SDXL_LIGHTNING_MODEL_ID = "ByteDance/SDXL-Lightning" SDXL_LIGHTNING_PATTERN = "sdxl_lightning_4step_lora.safetensors" -OPENAI_STRENGTH = 0.10 -GEMINI_STRENGTH = 0.15 -UNKNOWN_STRENGTH = GEMINI_STRENGTH -DEFAULT_STRENGTH = UNKNOWN_STRENGTH - -QWEN_OPENAI_STRENGTH = 0.10 -QWEN_GEMINI_STRENGTH = 0.25 -QWEN_UNKNOWN_STRENGTH = QWEN_GEMINI_STRENGTH +# Both profiles run the same distilled four-step schedule, and both are certified at a +# fixed seed because SynthID removal near the strength floor is seed-dependent. +PROFILE_STEPS = 4 +PROFILE_SEED = 0 +PROFILE_CFG = 1.0 # 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 @@ -55,74 +61,64 @@ class _StrengthPolicy: return self.by_vendor.get((vendor or "").casefold(), self.unknown) -_STANDARD_POLICY = _StrengthPolicy( - unknown=UNKNOWN_STRENGTH, - by_vendor={"openai": OPENAI_STRENGTH, "google": GEMINI_STRENGTH}, -) -_QWEN_POLICY = _StrengthPolicy( - unknown=QWEN_UNKNOWN_STRENGTH, - by_vendor={"openai": QWEN_OPENAI_STRENGTH, "google": QWEN_GEMINI_STRENGTH}, -) _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: - """Normalize spelling and resolve compatibility aliases.""" + """Normalize spelling and resolve the underscore spellings.""" value = profile.strip().casefold() return _ALIASES.get(value, value) -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) in _FOUR_STEP_PROFILES else 50 +def resolve_steps(num_inference_steps: int | None) -> int: + """Return an explicit step count or the distilled four-step default.""" + return PROFILE_STEPS if num_inference_steps is None else num_inference_steps -def resolve_seed(seed: int | None, pipeline: str) -> int | None: - """Keep the fixed four-step Z-Image profiles reproducible by default.""" - if seed is not None: - return seed - return 0 if normalize_profile(pipeline) in _FOUR_STEP_PROFILES else None +def resolve_seed(seed: int | None) -> int: + """Keep both profiles reproducible by default.""" + return PROFILE_SEED if seed is None else seed def strength_default_help() -> str: """Describe the live default policy without duplicating its values.""" return ( - f"vendor-adaptive (OpenAI {OPENAI_STRENGTH} / Google {GEMINI_STRENGTH} / " - f"unknown {UNKNOWN_STRENGTH}, from the C2PA issuer; qwen-zimage instead uses " - "resolution-adaptive denoise)" + "profile-adaptive (qwen-zimage uses resolution-adaptive denoise; sdxl-zimage " + f"uses OpenAI {SDXL_ZIMAGE_OPENAI_STRENGTH} / Google {SDXL_ZIMAGE_GEMINI_STRENGTH} / " + f"unknown {SDXL_ZIMAGE_UNKNOWN_STRENGTH}, from the C2PA issuer)" ) -def resolve_strength(strength: float | None, vendor: str | None = None, pipeline: str | None = None) -> float: - """Resolve a user override or the calibrated policy for a profile and vendor.""" +def resolve_strength( + strength: float | None, + vendor: str | None = None, + pipeline: str | None = None, + *, + size: tuple[int, int] | None = None, +) -> float: + """Resolve a user override or the calibrated policy for a profile and vendor. + + Total by design. qwen-zimage picks its strength from image area rather than from + the vendor, so it needs ``size``; returning ``None`` for it instead would push that + branch onto every caller and move one of the two strength policies outside this + module. ``size`` is required for qwen-zimage without an explicit strength. + """ if strength is not None: return strength - 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) + if normalize_profile(pipeline or "") == SDXL_ZIMAGE_PROFILE: + return _SDXL_ZIMAGE_POLICY.choose(vendor) + if size is None: + raise ValueError("qwen-zimage resolves strength from image area, so size is required") + from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise - -def viable_steps(num_inference_steps: int, strength: float) -> int: - """Ensure Diffusers receives at least one effective img2img denoising step.""" - if strength <= 0 or int(num_inference_steps * strength) >= 1: - return num_inference_steps - return math.ceil(1.0 / strength) + return resolution_adaptive_denoise(*size) def vendor_for_strength(image_path: Path) -> Literal["openai", "google"] | None: diff --git a/src/remove_ai_watermarks/_internal/watermark_remover.py b/src/remove_ai_watermarks/_internal/watermark_remover.py index 9dc6067..0dc4bb1 100644 --- a/src/remove_ai_watermarks/_internal/watermark_remover.py +++ b/src/remove_ai_watermarks/_internal/watermark_remover.py @@ -12,18 +12,19 @@ from typing import TYPE_CHECKING, Any from PIL import Image from remove_ai_watermarks._internal.watermark_profiles import ( - CONTROLNET_CANNY_MODEL, DEFAULT_MODEL_ID, - DEFAULT_STRENGTH, - QWEN_MODEL_ID, + DEFAULT_PROFILE, + PROFILE_CFG, + PROFILE_CHOICES, + PROFILE_STEPS, QWEN_ZIMAGE_PROFILE, SDXL_ZIMAGE_PROFILE, normalize_profile, resolve_seed, resolve_steps, resolve_strength, - viable_steps, ) +from remove_ai_watermarks.optional_deps import module_available if TYPE_CHECKING: from collections.abc import Callable @@ -33,7 +34,6 @@ 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", @@ -47,22 +47,7 @@ except ImportError: torch = None # type: ignore[assignment] _HAS_TORCH = False -try: - from diffusers import AutoPipelineForImage2Image as AutoImg2ImgPipeline - - _HAS_DIFFUSERS = True -except ImportError: - AutoImg2ImgPipeline = None # type: ignore[assignment,misc] - _HAS_DIFFUSERS = False - -_SDXL_FP16_VAE_ID = "madebyollin/sdxl-vae-fp16-fix" -_DEGENERATE_THRESHOLD = 1.0 -_CANNY_LOW = 100 -_CANNY_HIGH = 200 -_CONTROLNET_PROMPT = "best quality, high quality, sharp, detailed, photographic" -_CONTROLNET_NEGATIVE = "blurry, lowres, deformed, distorted text, garbled text, watermark, jpeg artifacts" -_QWEN_PROMPT = "high quality, sharp, detailed, faithful to the original" -_QWEN_NEGATIVE = "blurry, lowres, distorted text, garbled text, artifacts" +_HAS_DIFFUSERS = module_available("diffusers") def is_watermark_removal_available() -> bool: @@ -77,19 +62,6 @@ def _ensure_watermark_deps() -> None: ) -def _needs_fp16_vae_fix(model_id: str, default_model_id: str, is_fp16: bool) -> bool: - """Return whether the default SDXL pipeline needs the overflow-safe VAE.""" - return is_fp16 and model_id == default_model_id - - -def _is_degenerate_image(image: Image.Image) -> bool: - """Detect the uniform near-black output produced by an fp16 decode collapse.""" - import numpy as np - - pixels = np.asarray(image.convert("RGB"), dtype=np.float32) - return float(pixels.mean()) < _DEGENERATE_THRESHOLD and float(pixels.std()) < _DEGENERATE_THRESHOLD - - def _has_nvidia_gpu() -> bool: try: subprocess.run( @@ -103,23 +75,21 @@ def _has_nvidia_gpu() -> bool: return True -def _detect_cuda_index_url() -> str: - """Return a PyTorch wheel index compatible with the reported CUDA runtime.""" - try: - report = subprocess.run( - ["nvidia-smi"], - check=True, - capture_output=True, - text=True, - ).stdout - except (FileNotFoundError, subprocess.CalledProcessError): - return "https://download.pytorch.org/whl/cu121" - import re +def try_empty_device_cache(device: str) -> None: + """Ask Torch to release cached accelerator memory when the backend supports it. - match = re.search(r"CUDA Version:\s*(\d+)\.(\d+)", report) - if match is None: - return "https://download.pytorch.org/whl/cu121" - return f"https://download.pytorch.org/whl/cu{match.group(1)}{match.group(2)}" + Moved here when ``img2img_runner`` was deleted: the runner and its MPS recovery + path went with the CPU/MPS profiles, leaving this as that module's only content. + Silent by design -- it runs in cleanup paths where a raise would replace the real + error. + """ + if not _HAS_TORCH: + return + backend = getattr(torch, device, None) # type: ignore[union-attr] + empty_cache = getattr(backend, "empty_cache", None) + if callable(empty_cache): + with contextlib.suppress(Exception): + empty_cache() def _backend_works(device: str) -> bool: @@ -148,48 +118,11 @@ def get_device() -> str: return "cpu" -def _make_seed_generator(device: str, seed: int) -> Any: - """Create a deterministic generator, using CPU when device RNG is unavailable.""" - try: - return torch.Generator(device=device).manual_seed(seed) # type: ignore[union-attr] - except (RuntimeError, TypeError): - return torch.Generator().manual_seed(seed) # type: ignore[union-attr] - - -def _qwen_target_size(width: int, height: int) -> tuple[int, int]: - """Floor dimensions to Qwen's 16-pixel latent grid.""" - return max(16, width - width % 16), max(16, height - height % 16) - - -def _build_qwen_kwargs( - image: Image.Image, - strength: float, - num_inference_steps: int, - true_cfg_scale: float, - generator: Any, -) -> dict[str, Any]: - """Build the Qwen img2img call without importing its optional pipeline class.""" - width, height = _qwen_target_size(image.width, image.height) - return { - "prompt": _QWEN_PROMPT, - "negative_prompt": _QWEN_NEGATIVE, - "image": image, - "strength": strength, - "num_inference_steps": num_inference_steps, - "true_cfg_scale": true_cfg_scale, - "generator": generator, - "width": width, - "height": height, - } - - class WatermarkRemover: """Load one regeneration profile and write a metadata-clean raster output.""" DEFAULT_MODEL_ID = DEFAULT_MODEL_ID - DEFAULT_STRENGTH = DEFAULT_STRENGTH - CONTROLNET_CANNY_MODEL = CONTROLNET_CANNY_MODEL - _DEVICES = frozenset({"cpu", "mps", "cuda", "xpu"}) + _DEVICES = frozenset({"cuda"}) def __init__( self, @@ -198,47 +131,48 @@ class WatermarkRemover: torch_dtype: Any = None, progress_callback: Callable[[str], None] | None = None, hf_token: str | None = None, - pipeline: str = "controlnet", + pipeline: str = DEFAULT_PROFILE, controlnet_conditioning_scale: float = 1.0, cpu_offload: bool = False, ) -> None: - requested_model = model_id or self.DEFAULT_MODEL_ID self.model_profile = normalize_profile(pipeline) - if self.model_profile in _ZIMAGE_PROFILES and model_id not in {None, self.DEFAULT_MODEL_ID}: + if self.model_profile not in PROFILE_CHOICES: + raise ValueError(f"Unsupported pipeline '{pipeline}'. Use one of: {', '.join(PROFILE_CHOICES)}.") + if model_id is not None: 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 - else requested_model + else f"{DEFAULT_MODEL_ID} + Tongyi-MAI/Z-Image-Turbo" ) _ensure_watermark_deps() selected_device = (device or get_device()).casefold() self.device = get_device() if selected_device == "auto" else selected_device + # CUDA is a precondition of the object, not of the run. Both profiles raise on + # any other device, so accepting one here only defers a guaranteed failure to + # model-load time, several layers down and under the wrong profile's name. if self.device not in self._DEVICES: - raise ValueError(f"Unsupported device '{device}'. Use one of: auto, cpu, mps, cuda, xpu.") + raise ValueError( + f"Invisible-watermark removal is CUDA-only, so '{device}' cannot run it. " + "Both remaining profiles need an NVIDIA GPU. Visible-mark removal and " + "every identify command still run on CPU." + ) if torch_dtype is not None: 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: - self.torch_dtype = torch.float16 # type: ignore[union-attr] + self.torch_dtype = torch.bfloat16 # type: ignore[union-attr] self.cpu_offload = cpu_offload self.controlnet_conditioning_scale = controlnet_conditioning_scale self.hf_token = hf_token or os.environ.get("HF_TOKEN") self._progress_callback = progress_callback - self._pipeline: Any = None - self._controlnet_pipeline: Any = None - self._qwen_pipeline: Any = None self._qwen_zimage_pipeline: Any = None def _set_progress(self, message: str) -> None: @@ -248,114 +182,7 @@ class WatermarkRemover: def preload(self, *, global_only: bool = False) -> None: """Materialize the selected model stack before the first request.""" - 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() - elif self.model_profile == "controlnet": - self._load_controlnet_pipeline() - else: - self._load_pipeline() - - def _base_load_kwargs(self) -> dict[str, Any]: - options: dict[str, Any] = {"torch_dtype": self.torch_dtype} - if self.hf_token: - options["token"] = self.hf_token - return options - - def _load_from_pretrained(self, cls: Any, model_id: str, **kwargs: Any) -> Any: - if self.torch_dtype == torch.float16: # type: ignore[union-attr] - try: - return cls.from_pretrained(model_id, variant="fp16", **kwargs) - except Exception as error: - logger.info("Model %s has no usable fp16 variant (%s); using default weights", model_id, error) - return cls.from_pretrained(model_id, **kwargs) - - def _maybe_add_fp16_vae(self, options: dict[str, Any]) -> None: - if not _needs_fp16_vae_fix( - self.model_id, - self.DEFAULT_MODEL_ID, - self.torch_dtype == torch.float16, # type: ignore[union-attr] - ): - return - from diffusers import AutoencoderKL - - options["vae"] = AutoencoderKL.from_pretrained(_SDXL_FP16_VAE_ID, torch_dtype=torch.float16) - - @staticmethod - def _disable_sdxl_watermarker(options: dict[str, Any]) -> None: - options["add_watermarker"] = False - - def _move_to_device_and_optimize(self, pipeline: Any) -> Any: - if self.cpu_offload and self.device == "cuda": - offload = getattr(pipeline, "enable_model_cpu_offload", None) - if not callable(offload): - raise RuntimeError("CPU offload was requested, but this pipeline does not support it.") - offload(device="cuda") - else: - try: - pipeline = pipeline.to(self.device) - except (RuntimeError, AssertionError) as error: - if self.device == "cuda": - raise RuntimeError( - f"Failed to move model to CUDA ({error}). Install a compatible PyTorch wheel from " - f"{_detect_cuda_index_url()}." - ) from error - raise - optimize = getattr(pipeline, "enable_xformers_memory_efficient_attention", None) - if callable(optimize): - with contextlib.suppress(Exception): - optimize() - if self.device == "mps": - slice_attention = getattr(pipeline, "enable_attention_slicing", None) - if callable(slice_attention): - with contextlib.suppress(Exception): - slice_attention("max") - return pipeline - - def _sdxl_options(self) -> dict[str, Any]: - options = self._base_load_kwargs() - self._disable_sdxl_watermarker(options) - self._maybe_add_fp16_vae(options) - return options - - def _load_pipeline(self) -> Any: - if self._pipeline is None: - options = self._sdxl_options() - options.update(safety_checker=None, requires_safety_checker=False) - loaded = self._load_from_pretrained(AutoImg2ImgPipeline, self.model_id, **options) - self._pipeline = self._move_to_device_and_optimize(loaded) - return self._pipeline - - def _load_controlnet_pipeline(self) -> Any: - if self._controlnet_pipeline is None: - from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline - - controlnet = self._load_from_pretrained( - ControlNetModel, - CONTROLNET_CANNY_MODEL, - torch_dtype=self.torch_dtype, - ) - options = self._sdxl_options() - options["controlnet"] = controlnet - loaded = self._load_from_pretrained( - StableDiffusionXLControlNetImg2ImgPipeline, - self.model_id, - **options, - ) - self._controlnet_pipeline = self._move_to_device_and_optimize(loaded) - return self._controlnet_pipeline - - def _load_qwen_pipeline(self) -> Any: - if self._qwen_pipeline is None: - try: - from diffusers import QwenImageImg2ImgPipeline - except ImportError as error: - raise ImportError("The qwen profile requires Diffusers with QwenImageImg2ImgPipeline.") from error - model_id = QWEN_MODEL_ID if self.model_id == self.DEFAULT_MODEL_ID else self.model_id - loaded = QwenImageImg2ImgPipeline.from_pretrained(model_id, **self._base_load_kwargs()) - self._qwen_pipeline = self._move_to_device_and_optimize(loaded) - return self._qwen_pipeline + self._load_qwen_zimage_pipeline().preload(global_only=global_only) def _load_qwen_zimage_pipeline(self) -> Any: if self._qwen_zimage_pipeline is None: @@ -379,88 +206,6 @@ class WatermarkRemover: ) return self._qwen_zimage_pipeline - def _reload_on_cpu(self, cache_name: str, loader: Callable[[], Any]) -> Any: - self.device = "cpu" - self.torch_dtype = torch.float32 # type: ignore[union-attr] - setattr(self, cache_name, None) - return loader() - - def _run_img2img( - self, - init_image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - ) -> Image.Image: - from remove_ai_watermarks._internal.img2img_runner import run_img2img_with_mps_fallback - - output, device = run_img2img_with_mps_fallback( - self._load_pipeline, - init_image, - strength, - num_inference_steps, - guidance_scale, - generator, - self.device, - self._set_progress, - reload_on_cpu=lambda: self._reload_on_cpu("_pipeline", self._load_pipeline), - ) - self.device = device - return output - - def _build_canny_control_image(self, init_image: Image.Image) -> Image.Image: - import cv2 - import numpy as np - - gray = cv2.cvtColor(np.asarray(init_image.convert("RGB")), cv2.COLOR_RGB2GRAY) - edges = cv2.Canny(gray, _CANNY_LOW, _CANNY_HIGH) - return Image.fromarray(np.repeat(edges[:, :, None], 3, axis=2)) - - def _run_controlnet( - self, - init_image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - ) -> Image.Image: - from remove_ai_watermarks._internal.img2img_runner import run_img2img_with_mps_fallback - - extras = { - "prompt": _CONTROLNET_PROMPT, - "negative_prompt": _CONTROLNET_NEGATIVE, - "control_image": self._build_canny_control_image(init_image), - "controlnet_conditioning_scale": float(self.controlnet_conditioning_scale), - } - output, device = run_img2img_with_mps_fallback( - self._load_controlnet_pipeline, - init_image, - strength, - num_inference_steps, - guidance_scale, - generator, - self.device, - self._set_progress, - reload_on_cpu=lambda: self._reload_on_cpu("_controlnet_pipeline", self._load_controlnet_pipeline), - extra_kwargs=extras, - ) - self.device = device - return output - - def _run_qwen( - self, - init_image: Image.Image, - strength: float, - num_inference_steps: int, - guidance_scale: float, - generator: Any, - ) -> Image.Image: - response = self._load_qwen_pipeline()( - **_build_qwen_kwargs(init_image, strength, num_inference_steps, guidance_scale, generator) - ) - return response.images[0] - def _run_qwen_zimage( self, init_image: Image.Image, @@ -484,40 +229,20 @@ class WatermarkRemover: self, image: Image.Image, strength: float, - steps: int, - guidance: float, - generator: Any, seed: int | None, *, tile: bool, tile_size: int, tile_overlap: int, ) -> Image.Image: - if self.model_profile in _ZIMAGE_PROFILES: - return self._run_qwen_zimage( - image, - strength, - seed, - tile=tile, - tile_size=tile_size, - tile_overlap=tile_overlap, - ) - - runner = { - "qwen": self._run_qwen, - "controlnet": self._run_controlnet, - }.get(self.model_profile, self._run_img2img) - if tile and max(image.size) > tile_size: - from remove_ai_watermarks._internal.tiling import run_tiled - - return run_tiled( - lambda crop: runner(crop, strength, steps, guidance, generator), - image, - tile_size, - tile_overlap, - self._set_progress, - ) - return runner(image, strength, steps, guidance, generator) + return self._run_qwen_zimage( + image, + strength, + seed, + tile=tile, + tile_size=tile_size, + tile_overlap=tile_overlap, + ) def _write_output(self, image: Image.Image, output_path: Path) -> None: import numpy as np @@ -554,56 +279,27 @@ class WatermarkRemover: with Image.open(image_path) as opened: source = opened.convert("RGB") - if self.model_profile == QWEN_ZIMAGE_PROFILE: - from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise - - resolved_strength = strength if strength is not None else resolution_adaptive_denoise(*source.size) - else: - resolved_strength = resolve_strength(strength, vendor, self.model_profile) + resolved_strength = resolve_strength(strength, vendor, self.model_profile, size=source.size) if not 0.0 <= resolved_strength <= 1.0: raise ValueError(f"Strength must be between 0.0 and 1.0, got {resolved_strength}") - 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 in _ZIMAGE_PROFILES else guidance_scale or 7.5 - if self.model_profile in _ZIMAGE_PROFILES: - if steps != 4: - raise ValueError(f"The {self.model_profile} profile requires 4 steps.") - if guidance != 1.0: - raise ValueError(f"The {self.model_profile} profile requires CFG 1.0.") - else: - steps = viable_steps(steps, resolved_strength) + # Both profiles are distilled four-step schedules at CFG 1.0. Anything else is + # a caller error rather than a knob, so it is rejected instead of coerced. + steps = resolve_steps(num_inference_steps) + if steps != PROFILE_STEPS: + raise ValueError(f"The {self.model_profile} profile requires {PROFILE_STEPS} steps.") + if guidance_scale is not None and guidance_scale != PROFILE_CFG: + raise ValueError(f"The {self.model_profile} profile requires CFG {PROFILE_CFG}.") - generator = None - if resolved_seed is not None and _HAS_TORCH: - generator = _make_seed_generator(self.device, resolved_seed) result = self._generate( source, resolved_strength, - steps, - guidance, - generator, - resolved_seed, + resolve_seed(seed), tile=tile, tile_size=tile_size, tile_overlap=tile_overlap, ) - if self.torch_dtype == torch.float16 and _is_degenerate_image(result): # type: ignore[union-attr] - self.torch_dtype = torch.float32 # type: ignore[union-attr] - self._pipeline = self._controlnet_pipeline = self._qwen_pipeline = self._qwen_zimage_pipeline = None - result = self._generate( - source, - resolved_strength, - steps, - guidance, - generator, - resolved_seed, - tile=tile, - tile_size=tile_size, - tile_overlap=tile_overlap, - ) - if region is not None: import numpy as np @@ -634,8 +330,6 @@ class WatermarkRemover: if not input_dir.exists(): raise FileNotFoundError(f"Input directory not found: {input_dir}") output_dir.mkdir(parents=True, exist_ok=True) - from remove_ai_watermarks._internal.img2img_runner import try_empty_device_cache - outputs: list[Path] = [] candidates = sorted(path for path in input_dir.iterdir() if path.suffix.casefold() in extensions) for source in candidates: diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index d106289..e55d627 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -24,6 +24,9 @@ from remove_ai_watermarks import __version__, image_io, watermark_registry from remove_ai_watermarks._internal.constants import SUPPORTED_FORMATS from remove_ai_watermarks._internal.utils import is_supported_format from remove_ai_watermarks._internal.watermark_profiles import ( + DEFAULT_PROFILE, + PROFILE_CHOICES, + QWEN_ZIMAGE_PROFILE, resolve_seed, resolve_steps, resolve_strength, @@ -155,15 +158,15 @@ def _resolved_strength_for_display( vendor: str | None, pipeline: str, ) -> float: - """Resolve the same profile-specific strength the engine will execute.""" - if pipeline == "qwen-zimage" and strength is None: - from PIL import Image + """Resolve the same profile-specific strength the engine will execute. - from remove_ai_watermarks._internal.qwen_zimage_pipeline import resolution_adaptive_denoise + One call for both profiles, so the printed value cannot drift from the executed + one; the size is what qwen-zimage derives its strength from. + """ + from PIL import Image - with Image.open(source) as image: - return resolution_adaptive_denoise(image.width, image.height) - return resolve_strength(strength, vendor, pipeline) + with Image.open(source) as image: + return resolve_strength(strength, vendor, pipeline, size=image.size) # Shared option decorator for commands that run the invisible-watermark pipeline. @@ -173,8 +176,8 @@ _controlnet_scale_option = click.option( "--controlnet-scale", type=float, default=1.0, - help="ControlNet conditioning scale (structure/text preservation strength); " - "applies to the controlnet pipeline (the default). Higher = closer to original structure.", + help="Canny ControlNet conditioning scale on the global stage " + "(structure/text preservation strength). Higher = closer to original structure.", ) _min_resolution_option = click.option( @@ -203,9 +206,9 @@ _auto_option = click.option( "--auto", is_flag=True, default=False, - help="DEPRECATED: controlnet and adaptive polish are already the defaults, so " - "--auto only emits a warning and changes nothing. Use --no-adaptive-polish " - "to disable polishing.", + help="DEPRECATED: it no longer selects a pipeline. It now only requests the " + "adaptive polish, which the two-stage profiles otherwise leave off to keep their " + "output untouched. Prefer --adaptive-polish.", ) _adaptive_polish_option = click.option( @@ -251,49 +254,29 @@ _model_option = click.option( "--model", type=str, default=None, - help="HuggingFace model ID for the diffusion pipeline. Default: the SDXL base checkpoint.", + help="HuggingFace model ID. Both profiles pin a fixed model stack, so anything " + "other than the default is rejected rather than silently ignored.", ) _guidance_scale_option = click.option( "--guidance-scale", type=float, default=None, - help="Classifier-free guidance scale (CFG). Default: 7.5, except qwen-zimage " - "fixes CFG at 1.0. Lower = follow the prompt less / stay closer to the input.", + help="Classifier-free guidance scale (CFG). Both profiles are distilled and fix " + "CFG at 1.0, so any other value is rejected.", ) -def _normalize_pipeline(ctx: click.Context, param: click.Parameter, value: str | None) -> str | None: - """Resolve the legacy ``default`` profile name to ``sdxl`` (click option callback). - - Emits a one-line deprecation notice when the user explicitly passes the outdated - ``default`` value, pointing at the two current choices (``sdxl`` / ``controlnet``). - """ - if value is None: - return None - from remove_ai_watermarks._internal.watermark_profiles import normalize_profile - - normalized = normalize_profile(value) - if value.strip().lower() == "default": - click.echo( - "Warning: --pipeline default is deprecated and maps to 'sdxl'. " - "Use --pipeline sdxl (plain SDXL) or --pipeline controlnet (the default).", - err=True, - ) - return normalized - - -# ``controlnet`` (the default-SELECTED value), ``sdxl`` (plain SDXL img2img) and -# ``qwen`` (Qwen-Image, CUDA/cloud-class) are the current profiles; ``default`` is an -# OUTDATED back-compat alias for ``sdxl`` (warned + normalized away by _normalize_pipeline). -_PIPELINE_CHOICES = ["sdxl", "controlnet", "qwen", "qwen-zimage", "default"] +# The two-stage profiles are the only ones left. The former controlnet, sdxl, qwen and +# default profiles were removed rather than kept as a CPU path: none matched this +# recipe's face preservation, so offering them implied a quality the library no longer +# delivers. BOTH remaining profiles are CUDA-only. +_PIPELINE_CHOICES = list(PROFILE_CHOICES) _PIPELINE_HELP = ( - "Pipeline profile. controlnet (DEFAULT) = SDXL + canny ControlNet that preserves " - "text/faces via edge conditioning while removing SynthID; sdxl = plain SDXL img2img " - "(lighter, no extra model download, but leaves SynthID on flat-graphic content); " - "qwen = Qwen-Image (20B, Apache-2.0) img2img, best text/structure preservation but " - "CUDA/cloud-class; qwen-zimage = Qwen-Image-2512 + Lightning + Canny, followed by " - "SAM-masked Z-Image face repair (CUDA-only; install the qwen-zimage extra). " - "('default' is an OUTDATED alias for 'sdxl'.)" + "Pipeline profile. qwen-zimage (DEFAULT) = Qwen-Image-2512 + Lightning + Canny, " + "followed by SAM-masked Z-Image face repair; sdxl-zimage = the same recipe and the " + "same face stage on an SDXL global pass, which needs more denoise. Both are " + "CUDA-ONLY -- install the qwen-zimage extra. There is no CPU or MPS profile for " + "invisible-watermark removal." ) # Shared --pipeline / --strength decorators so the three diffusion commands @@ -302,8 +285,7 @@ _PIPELINE_HELP = ( _pipeline_option = click.option( "--pipeline", type=click.Choice(_PIPELINE_CHOICES), - default="controlnet", - callback=_normalize_pipeline, + default=DEFAULT_PROFILE, help=_PIPELINE_HELP, ) _strength_option = click.option( @@ -362,12 +344,10 @@ _visible_sensitivity_option = click.option( def _resolve_auto_polish(auto: bool, adaptive_polish: bool) -> bool: """Warn on the retired ``--auto`` flag, returning ``adaptive_polish`` unchanged. - ``--auto`` used to plan the pipeline + polish from content detection, but the - pipeline is now always controlnet (the default) and the adaptive polish is ON by - default (it self-gates by detail level), so the content detectors were removed and - ``--auto`` is now a no-op alias: the polish it used to enable is already the default, - and an explicit ``--no-adaptive-polish`` still wins. So it only emits a deprecation - warning and passes ``adaptive_polish`` through. + ``--auto`` used to plan the pipeline + polish from content detection. There is now + only one default pipeline, and the content detectors were removed, so the flag + survives purely as a polish request: it emits a deprecation warning and passes + ``adaptive_polish`` through, with an explicit ``--no-adaptive-polish`` still winning. """ if auto: click.echo( @@ -379,9 +359,14 @@ def _resolve_auto_polish(auto: bool, adaptive_polish: bool) -> bool: def _resolve_profile_polish(auto: bool, adaptive_polish: bool, pipeline: str) -> bool: - """Keep the upstream qwen-zimage output unchanged unless polish was explicit.""" + """Keep the upstream qwen-zimage output unchanged unless polish was explicit. + + ``--auto`` counts as explicit. It is deprecated, but it is still a request for the + polish, and once qwen-zimage became the DEFAULT pipeline the source check below + would otherwise have silently turned that flag into a no-op for every caller. + """ adaptive_polish = _resolve_auto_polish(auto, adaptive_polish) - if pipeline != "qwen-zimage": + if pipeline != QWEN_ZIMAGE_PROFILE or auto: return adaptive_polish ctx = click.get_current_context(silent=True) if ctx is None: @@ -885,7 +870,7 @@ def cmd_erase( "--steps", type=int, default=None, - help="Number of denoising steps. Default: 4 for qwen-zimage, 50 otherwise.", + help="Number of denoising steps. Both profiles are distilled four-step schedules, so 4 is the only accepted value.", ) @_pipeline_option @click.option( @@ -965,8 +950,8 @@ def cmd_invisible( from remove_ai_watermarks.invisible_engine import InvisibleEngine source = _validate_image(source) - steps = resolve_steps(steps, pipeline) - seed = resolve_seed(seed, pipeline) + steps = resolve_steps(steps) + seed = resolve_seed(seed) _warn_if_esrgan_unavailable(upscaler) adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline) if output is None: @@ -1575,7 +1560,7 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo "--steps", type=int, default=None, - help="Number of denoising steps. Default: 4 for qwen-zimage, 50 otherwise.", + help="Number of denoising steps. Both profiles are distilled four-step schedules, so 4 is the only accepted value.", ) @_pipeline_option @_model_option @@ -1651,8 +1636,8 @@ def cmd_all( """ _banner() source = _validate_image(source) - steps = resolve_steps(steps, pipeline) - seed = resolve_seed(seed, pipeline) + steps = resolve_steps(steps) + seed = resolve_seed(seed) _warn_if_esrgan_unavailable(upscaler) adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline) @@ -2018,7 +2003,7 @@ def _process_batch_image( "--steps", type=int, default=None, - help="Number of denoising steps. Default: 4 for qwen-zimage, 50 otherwise.", + help="Number of denoising steps. Both profiles are distilled four-step schedules, so 4 is the only accepted value.", ) @_visible_backend_option @_visible_sensitivity_option @@ -2105,8 +2090,8 @@ def cmd_batch( if mode in ("invisible", "all"): _warn_if_esrgan_unavailable(upscaler) adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline) - steps = resolve_steps(steps, pipeline) - seed = resolve_seed(seed, pipeline) + steps = resolve_steps(steps) + seed = resolve_seed(seed) options = _BatchOptions( strength=strength, steps=steps, diff --git a/src/remove_ai_watermarks/invisible_engine.py b/src/remove_ai_watermarks/invisible_engine.py index b7d34fa..b655763 100644 --- a/src/remove_ai_watermarks/invisible_engine.py +++ b/src/remove_ai_watermarks/invisible_engine.py @@ -20,7 +20,9 @@ from ._internal.watermark_profiles import ( DEFAULT_MODEL_ID as DEFAULT_SDXL_MODEL_ID, ) from ._internal.watermark_profiles import ( + DEFAULT_PROFILE, resolve_seed, + resolve_steps, ) if TYPE_CHECKING: @@ -90,7 +92,7 @@ class InvisibleEngine: self, model_id: str | None = None, device: str | None = None, - pipeline: str = "controlnet", + pipeline: str = DEFAULT_PROFILE, hf_token: str | None = None, progress_callback: Callable[[str], None] | None = None, controlnet_conditioning_scale: float = 1.0, @@ -101,19 +103,16 @@ class InvisibleEngine: Args: model_id: HuggingFace model ID. None = use the SDXL base default. device: Device for inference (auto/cpu/mps/cuda/xpu). None = auto. - pipeline: Pipeline profile. "controlnet" (DEFAULT; SDXL + canny ControlNet - that preserves text/face structure via edge conditioning while removing - SynthID), "sdxl" (plain SDXL img2img, lighter but leaves SynthID on - 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), 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". + pipeline: Pipeline profile, one of "qwen-zimage" (DEFAULT; + Qwen-Image-2512 Lightning + Canny, then SAM-masked Z-Image face repair) + 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). BOTH ARE CUDA-ONLY -- there is no CPU or MPS path + for invisible-watermark removal. hf_token: HuggingFace API token. progress_callback: Optional callback for progress messages. - controlnet_conditioning_scale: ControlNet structure-preservation - strength (controlnet pipeline only). + controlnet_conditioning_scale: Canny ControlNet structure-preservation + strength on the global stage of both profiles. cpu_offload: Offload model components to CPU between CUDA calls instead of keeping the whole pipeline in VRAM, at the cost of speed. For qwen-zimage, force the face stack to offload instead of using automatic @@ -238,11 +237,8 @@ class InvisibleEngine: """ import tempfile - if num_inference_steps is None: - profile = getattr(self._remover, "model_profile", None) - 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) + num_inference_steps = resolve_steps(num_inference_steps) + seed = resolve_seed(seed) from PIL import Image, ImageOps diff --git a/tests/test_cli.py b/tests/test_cli.py index 95ca22d..37b31b5 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -327,7 +327,7 @@ class TestInvisibleCommand: expected = sample_png.with_stem(sample_png.stem + "_clean") assert expected.exists() - def test_invisible_adaptive_polish_on_by_default(self, runner, sample_png): + def test_invisible_adaptive_polish_off_by_default_under_qwen_zimage(self, runner, sample_png): mock_cls, mock_engine = _mock_invisible_engine() with ( patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True), @@ -336,8 +336,10 @@ class TestInvisibleCommand: ): result = runner.invoke(main, ["invisible", str(sample_png), "--force"]) assert result.exit_code == 0, result.output - # adaptive_polish is ON by default (self-gating, so a no-op where not needed). - assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is True + # The default profile is qwen-zimage, and _resolve_profile_polish keeps its + # output untouched unless polish was asked for explicitly. It stays available: + # passing --adaptive-polish still turns it on (covered separately). + assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is False # Default model is None (the SDXL base) and CFG is None (the library's 7.5). assert mock_cls.call_args.kwargs["model_id"] is None assert mock_engine.remove_watermark.call_args.kwargs["guidance_scale"] is None @@ -368,30 +370,16 @@ class TestInvisibleCommand: assert mock_cls.call_args.kwargs["model_id"] == "org/custom-sdxl" assert mock_engine.remove_watermark.call_args.kwargs["guidance_scale"] == 5.5 - def test_pipeline_default_alias_warns_and_maps_to_sdxl(self, runner, sample_png): - mock_cls, _mock_engine = _mock_invisible_engine() - with ( - patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True), - patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), - patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls), - ): - result = runner.invoke(main, ["invisible", str(sample_png), "--pipeline", "default", "--force"]) - assert result.exit_code == 0, result.output - # The legacy value warns and is normalized to "sdxl" before the engine is built. - assert "deprecated" in result.output.lower() - assert mock_cls.call_args.kwargs["pipeline"] == "sdxl" + def test_retired_pipeline_names_are_rejected_not_silently_remapped(self, runner, sample_png): + """default/sdxl/controlnet/qwen were removed with their CPU code paths. - def test_pipeline_sdxl_does_not_warn(self, runner, sample_png): - mock_cls, _mock_engine = _mock_invisible_engine() - with ( - patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True), - patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), - patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls), - ): - result = runner.invoke(main, ["invisible", str(sample_png), "--pipeline", "sdxl", "--force"]) - assert result.exit_code == 0, result.output - assert "deprecated" not in result.output.lower() - assert mock_cls.call_args.kwargs["pipeline"] == "sdxl" + Click rejects them at parse time. Mapping them onward would run a profile the + caller never chose, at a different strength and a different quality. + """ + for retired in ("default", "sdxl", "controlnet", "qwen"): + result = runner.invoke(main, ["invisible", str(sample_png), "--pipeline", retired, "--force"]) + assert result.exit_code == 2, result.output + assert "is not one of" in result.output def test_invisible_nonexistent_file(self, runner): result = runner.invoke(main, ["invisible", "/nonexistent/file.png"]) @@ -871,8 +859,10 @@ class TestBatchCommand: assert out[100, 100, 3] == 255 def test_batch_auto_is_deprecated_and_enables_polish(self, runner, tmp_path): - """--auto is retired: it warns and just enables the adaptive polish (the - pipeline is always the default controlnet now).""" + """--auto is retired: it warns and just enables the adaptive polish. + + It no longer selects a pipeline: qwen-zimage is the only default there is. + """ input_dir = _make_batch_dir(tmp_path, count=2) output_dir = tmp_path / "output" mock_cls, mock_engine = _mock_invisible_engine() @@ -890,7 +880,7 @@ class TestBatchCommand: assert "2 processed" in result.output assert "deprecated" in result.output.lower() # Pipeline stays the default controlnet; --auto only turned the polish on. - assert mock_cls.call_args.kwargs["pipeline"] == "controlnet" + assert mock_cls.call_args.kwargs["pipeline"] == "qwen-zimage" assert mock_engine.remove_watermark.call_args.kwargs["adaptive_polish"] is True def test_batch_default_output_dir(self, runner, tmp_path): diff --git a/tests/test_cpu_offload.py b/tests/test_cpu_offload.py index 8a50b49..b9fc71a 100644 --- a/tests/test_cpu_offload.py +++ b/tests/test_cpu_offload.py @@ -8,10 +8,6 @@ core CI matrix needs no diffusion dependency, model download, or GPU. from __future__ import annotations -from unittest.mock import Mock - -import pytest - from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover @@ -21,55 +17,3 @@ def _remover(device: str, cpu_offload: bool) -> WatermarkRemover: remover.cpu_offload = cpu_offload remover._progress_callback = None return remover - - -class TestCpuOffloadPlacement: - def test_offload_enabled_on_cuda_streams_instead_of_moving(self): - remover = _remover("cuda", cpu_offload=True) - pipeline = Mock() - - returned = remover._move_to_device_and_optimize(pipeline) - - pipeline.enable_model_cpu_offload.assert_called_once_with(device="cuda") - pipeline.to.assert_not_called() - # Offload leaves the pipeline object in place (accelerate hooks handle it). - assert returned is pipeline - - def test_no_offload_moves_whole_pipeline_to_cuda(self): - remover = _remover("cuda", cpu_offload=False) - pipeline = Mock() - - remover._move_to_device_and_optimize(pipeline) - - pipeline.to.assert_called_once_with("cuda") - pipeline.enable_model_cpu_offload.assert_not_called() - - def test_offload_flag_ignored_off_cuda(self): - # The flag is CUDA-only: on cpu it must still be a plain .to("cpu"). - remover = _remover("cpu", cpu_offload=True) - pipeline = Mock() - - remover._move_to_device_and_optimize(pipeline) - - pipeline.to.assert_called_once_with("cpu") - pipeline.enable_model_cpu_offload.assert_not_called() - - def test_offload_fails_loudly_when_pipeline_lacks_support(self): - remover = _remover("cuda", cpu_offload=True) - pipeline = Mock(spec=["to"]) - - with pytest.raises(RuntimeError, match="does not support"): - remover._move_to_device_and_optimize(pipeline) - - pipeline.to.assert_not_called() - - def test_qwen_zimage_forces_face_stack_offload(self): - remover = _remover("cuda", cpu_offload=True) - remover.torch_dtype = object() - remover.hf_token = None - remover.controlnet_conditioning_scale = 1.0 - remover._qwen_zimage_pipeline = None - - runtime = remover._load_qwen_zimage_pipeline() - - assert runtime.keep_face_models_on_device is False diff --git a/tests/test_img2img_runner.py b/tests/test_img2img_runner.py deleted file mode 100644 index 7adeea2..0000000 --- a/tests/test_img2img_runner.py +++ /dev/null @@ -1,133 +0,0 @@ -"""Unit tests for the MPS->CPU fallback orchestration (no GPU/model required). - -``img2img_runner`` has no torch import at module top -- the pipeline is -injected as a plain callable -- so the fallback control flow is fully -mockable. This guards the exact behavior hit in production on Apple Silicon: -a native-resolution SDXL run that OOMs on MPS must transparently retry on CPU, -while any non-MPS error must propagate unchanged. -""" - -from __future__ import annotations - -from unittest.mock import Mock - -import pytest - -from remove_ai_watermarks._internal import img2img_runner -from remove_ai_watermarks._internal.img2img_runner import ( - run_img2img, - run_img2img_with_mps_fallback, -) - -_MPS_OOM = "MPS backend out of memory (MPS allocated: 17.21 GiB, max allowed: 20.13 GiB)" - - -def _result(image: object) -> Mock: - """A stand-in for a diffusers pipeline output object (has .images).""" - out = Mock() - out.images = [image] - return out - - -class TestMpsFallback: - def test_mps_error_reloads_on_cpu_and_retries(self, monkeypatch: pytest.MonkeyPatch): - sentinel = object() - inner = Mock(side_effect=[RuntimeError(_MPS_OOM), sentinel]) - monkeypatch.setattr(img2img_runner, "run_img2img", inner) - load_pipeline = Mock(return_value="gpu_pipe") - reload_on_cpu = Mock(return_value="cpu_pipe") - - img, device = run_img2img_with_mps_fallback( - load_pipeline, object(), 0.05, 50, 7.5, "gen", "mps", lambda _m: None, reload_on_cpu=reload_on_cpu - ) - - assert (img, device) == (sentinel, "cpu") - reload_on_cpu.assert_called_once() - assert inner.call_count == 2 - # Retry must use the reloaded CPU pipeline, device "cpu", and drop the - # MPS generator (generator=None) so CPU runs deterministically. - retry_args = inner.call_args_list[1].args - assert retry_args[0] == "cpu_pipe" - assert retry_args[5] is None # generator - assert retry_args[6] == "cpu" # device - - def test_happy_path_returns_original_device_without_reload(self, monkeypatch: pytest.MonkeyPatch): - sentinel = object() - monkeypatch.setattr(img2img_runner, "run_img2img", Mock(return_value=sentinel)) - reload_on_cpu = Mock() - - img, device = run_img2img_with_mps_fallback( - Mock(return_value="gpu_pipe"), - object(), - 0.05, - 50, - 7.5, - "gen", - "mps", - lambda _m: None, - reload_on_cpu=reload_on_cpu, - ) - - assert (img, device) == (sentinel, "mps") - reload_on_cpu.assert_not_called() - - def test_non_mps_runtime_error_propagates(self, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(img2img_runner, "run_img2img", Mock(side_effect=RuntimeError("CUDA out of memory"))) - reload_on_cpu = Mock() - - with pytest.raises(RuntimeError, match="CUDA"): - run_img2img_with_mps_fallback( - Mock(return_value="gpu_pipe"), - object(), - 0.05, - 50, - 7.5, - "gen", - "mps", - lambda _m: None, - reload_on_cpu=reload_on_cpu, - ) - reload_on_cpu.assert_not_called() - - def test_mps_error_on_non_mps_device_propagates(self, monkeypatch: pytest.MonkeyPatch): - # An "mps"-worded error while running on cpu must NOT trigger the reload. - monkeypatch.setattr(img2img_runner, "run_img2img", Mock(side_effect=RuntimeError(_MPS_OOM))) - reload_on_cpu = Mock() - - with pytest.raises(RuntimeError, match="MPS backend"): - run_img2img_with_mps_fallback( - Mock(return_value="cpu_pipe"), - object(), - 0.05, - 50, - 7.5, - None, - "cpu", - lambda _m: None, - reload_on_cpu=reload_on_cpu, - ) - reload_on_cpu.assert_not_called() - - -class TestRunImg2Img: - def test_returns_first_image_from_pipeline_result(self): - sentinel = object() - pipeline = Mock(return_value=_result(sentinel)) - - out = run_img2img(pipeline, object(), 0.05, 50, 7.5, None, "cpu", lambda _m: None) - - assert out is sentinel - - def test_typeerror_on_callback_retries_without_callback(self): - # Older diffusers reject the progress callback kwarg with TypeError; - # run_img2img must retry once without it rather than fail. - sentinel = object() - pipeline = Mock(side_effect=[TypeError("unexpected keyword 'callback'"), _result(sentinel)]) - - out = run_img2img(pipeline, object(), 0.05, 50, 7.5, None, "cpu", lambda _m: None) - - assert out is sentinel - assert pipeline.call_count == 2 - # First attempt passes the progress callback; the retry omits it. - assert "callback" in pipeline.call_args_list[0].kwargs - assert "callback" not in pipeline.call_args_list[1].kwargs diff --git a/tests/test_invisible_engine.py b/tests/test_invisible_engine.py index 077bcbe..05d28ce 100644 --- a/tests/test_invisible_engine.py +++ b/tests/test_invisible_engine.py @@ -4,7 +4,6 @@ from __future__ import annotations from types import SimpleNamespace -import pytest from PIL import Image from remove_ai_watermarks.invisible_engine import InvisibleEngine, _target_size, is_available @@ -204,30 +203,3 @@ class TestEsrganUpscale: out = InvisibleEngine._esrgan_upscale(self._fake_engine(), img, (512, 341)) assert out.size == (512, 341) assert np.array_equal(np.asarray(out), np.asarray(img.resize((512, 341), Image.Resampling.LANCZOS))) - - -class TestCannyControlImage: - """The ControlNet canny conditioning image builder (pure cv2/numpy; behind the gpu - extra since it lives on WatermarkRemover). Skips when torch/diffusers are absent.""" - - def test_edge_map_is_3channel_rgb(self): - if not is_available(): - pytest.skip("diffusion extra (torch/diffusers) not installed") - import numpy as np - - from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover - - rng = np.random.default_rng(0) - img = Image.fromarray(rng.integers(0, 256, (64, 80, 3), dtype=np.uint8)) - # The method uses no instance state, so call it unbound with a dummy self. - out = WatermarkRemover._build_canny_control_image(None, img) # type: ignore[arg-type] - arr = np.array(out) - assert out.mode == "RGB" - assert arr.shape == (64, 80, 3) - import cv2 - - gray = cv2.cvtColor(np.asarray(img.convert("RGB")), cv2.COLOR_RGB2GRAY) - expected = cv2.Canny(gray, 100, 200) - assert np.array_equal(arr[:, :, 0], expected) - assert np.array_equal(arr[:, :, 1], expected) - assert np.array_equal(arr[:, :, 2], expected) diff --git a/tests/test_platform.py b/tests/test_platform.py index 694f03b..b2af277 100644 --- a/tests/test_platform.py +++ b/tests/test_platform.py @@ -9,17 +9,14 @@ from __future__ import annotations from pathlib import Path from unittest.mock import MagicMock, patch -import numpy as np import pytest -from PIL import Image -from remove_ai_watermarks._internal.progress import is_mps_error from remove_ai_watermarks._internal.utils import get_image_format, is_supported_format from remove_ai_watermarks._internal.watermark_profiles import ( - DEFAULT_STRENGTH, - GEMINI_STRENGTH, - OPENAI_STRENGTH, - UNKNOWN_STRENGTH, + PROFILE_CHOICES, + SDXL_ZIMAGE_GEMINI_STRENGTH, + SDXL_ZIMAGE_OPENAI_STRENGTH, + SDXL_ZIMAGE_UNKNOWN_STRENGTH, normalize_profile, resolve_strength, strength_default_help, @@ -59,290 +56,158 @@ class TestDeviceDetection: assert get_device() == "xpu" fake_torch.tensor.assert_called_with([1.0], device="xpu") - def test_init_accepts_xpu_and_selects_fp16(self): - """WatermarkRemover accepts device='xpu' and picks fp16 (not fp32).""" + def test_non_cuda_devices_are_refused_at_construction(self): + """CUDA is a precondition of the object, not of the run. + + Both remaining profiles raise on any other device, so accepting cpu/mps/xpu + here only defers a guaranteed failure to model-load time - several layers down, + after the dependency check and the pipeline import, under a message naming + whichever profile the internal pipeline happens to be. + """ if not is_watermark_removal_available(): pytest.skip("torch/diffusers not installed") import torch from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover - remover = WatermarkRemover(device="xpu") - assert remover.device == "xpu" - assert remover.torch_dtype == torch.float16 + for device in ("cpu", "mps", "xpu"): + with pytest.raises(ValueError, match="CUDA-only"): + WatermarkRemover(device=device) - def test_seed_generator_falls_back_to_cpu_when_device_rng_unsupported(self): - """A device with no RNG backend (e.g. some torch-xpu builds) falls back - to a CPU generator instead of raising when --seed is used.""" - from remove_ai_watermarks._internal import watermark_remover as wr - - def fake_generator(device="cpu"): - if device == "xpu": - raise RuntimeError("Device type xpu is not supported for torch.Generator()") - gen = MagicMock() - gen.manual_seed.return_value = f"gen:{device}" - return gen - - fake_torch = MagicMock() - fake_torch.Generator.side_effect = fake_generator - with patch.object(wr, "torch", fake_torch): - assert wr._make_seed_generator("xpu", 123) == "gen:cpu" - assert wr._make_seed_generator("cuda", 123) == "gen:cuda" + remover = WatermarkRemover(device="cuda") + assert remover.device == "cuda" + assert remover.torch_dtype == torch.bfloat16 -class TestMpsErrorDetection: - """Tests for MPS error detection helper.""" +class TestEmptyDeviceCache: + """try_empty_device_cache is all that remains of the img2img runner. - def test_detects_mps_error(self): - err = RuntimeError("MPS backend out of memory") - assert is_mps_error(err) is True + Its module lost run_img2img and the MPS fallback along with the CPU/MPS profiles; + both surviving profiles are CUDA-only, so there is no MPS failure left to recover + from. The helper must stay silent on a backend that cannot empty a cache, because + it runs in cleanup paths where a raise would replace the real error. + """ - def test_non_mps_error(self): - err = RuntimeError("CUDA out of memory") - assert is_mps_error(err) is False + def test_unknown_backend_is_a_silent_no_op(self): + from remove_ai_watermarks._internal.watermark_remover import try_empty_device_cache - def test_generic_error(self): - err = RuntimeError("something went wrong") - assert is_mps_error(err) is False - - -# ── Model profiles ────────────────────────────────────────────────── + try_empty_device_cache("cpu") + try_empty_device_cache("definitely-not-a-backend") class TestModelProfiles: - """Tests for watermark_profiles.py profile-name normalization.""" + """Only the two CUDA-only two-stage profiles remain.""" def test_canonical_profiles_unchanged(self): - assert normalize_profile("sdxl") == "sdxl" - assert normalize_profile("controlnet") == "controlnet" - assert normalize_profile("qwen") == "qwen" + assert normalize_profile("qwen-zimage") == "qwen-zimage" + assert normalize_profile("sdxl-zimage") == "sdxl-zimage" - def test_default_alias_resolves_to_sdxl(self): - # "default" is the legacy alias for "sdxl" (back-compat for existing scripts). - assert normalize_profile("default") == "sdxl" + def test_underscore_spellings_resolve(self): + assert normalize_profile("qwen_zimage") == "qwen-zimage" + assert normalize_profile(" SDXL_ZImage ") == "sdxl-zimage" - def test_normalize_is_case_and_whitespace_insensitive(self): - assert normalize_profile(" Default ") == "sdxl" - assert normalize_profile("CONTROLNET") == "controlnet" + def test_retired_names_no_longer_resolve_to_a_profile(self): + """default/sdxl/controlnet/qwen were removed, not aliased onward. - -class TestFp16WeightVariant: - """_load_from_pretrained reads the fp16 weight variant on fp16, with a fallback. - - Loading the fp16 ``variant`` reads the half-precision weight files (~half the bytes) - instead of the fp32 defaults + a downcast, which roughly halves the cold-start weight - read. fp32 (cpu/mps) and bf16 (qwen) must never request the variant; a checkpoint - without fp16 files must fall back to the default weights (prior behavior). - """ - - def _remover(self, dtype: object): - if not is_watermark_removal_available(): - pytest.skip("torch/diffusers not installed") - from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover - - # device="cpu" alone would force fp32; the explicit torch_dtype override lets us - # exercise the fp16 path with no GPU (construction loads no weights). - return WatermarkRemover(device="cpu", torch_dtype=dtype) - - def test_fp16_requests_variant(self): - import torch - - remover = self._remover(torch.float16) - cls = MagicMock() - cls.from_pretrained.return_value = "PIPE" - out = remover._load_from_pretrained(cls, "some/model", token="t") - assert out == "PIPE" - cls.from_pretrained.assert_called_once_with("some/model", variant="fp16", token="t") - - def test_fp16_falls_back_when_variant_missing(self): - import torch - - remover = self._remover(torch.float16) - cls = MagicMock() - cls.from_pretrained.side_effect = [OSError("no fp16 weight files"), "PIPE"] - out = remover._load_from_pretrained(cls, "some/model", token="t") - assert out == "PIPE" - assert cls.from_pretrained.call_count == 2 - first, second = cls.from_pretrained.call_args_list - assert first.kwargs.get("variant") == "fp16" - assert "variant" not in second.kwargs # the fallback drops the variant - - def test_fp32_never_requests_variant(self): - import torch - - remover = self._remover(torch.float32) - cls = MagicMock() - cls.from_pretrained.return_value = "PIPE" - remover._load_from_pretrained(cls, "some/model") - cls.from_pretrained.assert_called_once_with("some/model") - assert "variant" not in cls.from_pretrained.call_args.kwargs + Silently mapping them at the alias layer would route an old script into a + profile it never asked for; the remover raises on the unknown name instead. + """ + for retired in ("default", "sdxl", "controlnet", "qwen"): + assert normalize_profile(retired) not in PROFILE_CHOICES class TestNoReembeddedWatermark: - """F2 regression: the SDXL removal pipelines must disable the diffusers default - invisible watermarker (``add_watermarker=False``). + """F2 regression: the SDXL global stage must disable the diffusers watermarker. diffusers stamps an open "Stable Diffusion XL" DWT-DCT watermark onto every SDXL output whenever ``invisible-watermark`` is installed. A watermark REMOVER that left - it on would replace one detectable AI watermark (SynthID) with another -- the cleaned - output re-reads as AI. The ControlNet sub-model load must NOT receive the kwarg (it - is not a pipeline and does not accept it). + it on would replace one detectable AI watermark (SynthID) with another -- the + cleaned output re-reads as AI. The ControlNet sub-model load must NOT receive the + kwarg, since it is not a pipeline and does not accept it. + + Only sdxl-zimage carries an SDXL pipeline now; qwen-zimage's global stage is + DiffSynth, which has no such watermarker. """ - def _remover(self, profile: str): + def test_sdxl_global_stage_disables_watermarker(self, monkeypatch: pytest.MonkeyPatch): if not is_watermark_removal_available(): pytest.skip("torch/diffusers not installed") - from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover + import diffusers - return WatermarkRemover(device="cpu", pipeline=profile) + from remove_ai_watermarks._internal.sdxl_zimage_pipeline import SdxlZImagePipeline - def _capture(self, monkeypatch, remover): - from remove_ai_watermarks._internal.watermark_remover import WatermarkRemover + calls: dict[str, dict] = {} - calls: list[tuple[str, dict]] = [] + def record(name): + def fake(*_args, **kwargs): + calls[name] = kwargs + return MagicMock() - def fake_load(self, cls, model_id, **kwargs): - calls.append((getattr(cls, "__name__", str(cls)), kwargs)) - return MagicMock() + return fake - monkeypatch.setattr(WatermarkRemover, "_load_from_pretrained", fake_load) - monkeypatch.setattr(WatermarkRemover, "_move_to_device_and_optimize", lambda self, p: p) - return calls + monkeypatch.setattr(diffusers.ControlNetModel, "from_pretrained", record("controlnet")) + monkeypatch.setattr(diffusers.AutoencoderKL, "from_pretrained", record("vae")) + monkeypatch.setattr(diffusers.StableDiffusionXLControlNetImg2ImgPipeline, "from_pretrained", record("pipeline")) + monkeypatch.setattr("huggingface_hub.hf_hub_download", lambda *a, **k: "lora.safetensors") + # from_config would otherwise resolve the mock's config as a repo id. + monkeypatch.setattr(diffusers.EulerDiscreteScheduler, "from_config", lambda *a, **k: MagicMock()) - def test_sdxl_pipeline_disables_watermarker(self, monkeypatch: pytest.MonkeyPatch): - remover = self._remover("sdxl") - calls = self._capture(monkeypatch, remover) - remover._load_pipeline() - assert any(kw.get("add_watermarker") is False for _, kw in calls), calls + pipeline = SdxlZImagePipeline(device="cuda", torch_dtype=None) + monkeypatch.setattr(type(pipeline), "_require_cuda", lambda self: None) + pipeline._load_sdxl() - def test_controlnet_pipeline_disables_watermarker(self, monkeypatch: pytest.MonkeyPatch): - remover = self._remover("controlnet") - calls = self._capture(monkeypatch, remover) - remover._load_controlnet_pipeline() - by_cls = dict(calls) - # the SDXL pipeline load disables the watermarker... - assert by_cls["StableDiffusionXLControlNetImg2ImgPipeline"].get("add_watermarker") is False - # ...but the ControlNet sub-model load must not carry the kwarg (it would error). - assert "add_watermarker" not in by_cls["ControlNetModel"] - - -class _StubImage: - """Minimal PIL.Image stand-in: just the ``width``/``height`` the pure helper reads.""" - - def __init__(self, width: int, height: int) -> None: - self.width = width - self.height = height - - -class TestQwenKwargs: - """_build_qwen_kwargs is pure (no torch); guards the Qwen-Image call shape. - - watermark_remover imports torch under a try/except, so the module (and this pure - helper) imports fine in the default+dev CI env where torch is absent. - """ - - def test_uses_true_cfg_not_guidance_scale(self): - from remove_ai_watermarks._internal.watermark_remover import _build_qwen_kwargs - - gen = object() - img = _StubImage(2816, 1536) - kwargs = _build_qwen_kwargs(img, strength=0.3, num_inference_steps=40, true_cfg_scale=4.0, generator=gen) - # Qwen uses true_cfg_scale, NOT SDXL's guidance_scale. - assert kwargs["true_cfg_scale"] == 4.0 - assert "guidance_scale" not in kwargs - # The scrub still comes from strength; image + generator pass through. - assert kwargs["strength"] == 0.3 - assert kwargs["image"] is img - assert kwargs["generator"] is gen - assert kwargs["prompt"] == "high quality, sharp, detailed, faithful to the original" - assert kwargs["negative_prompt"] == "blurry, lowres, distorted text, garbled text, artifacts" - - def test_passes_explicit_aspect_preserving_size(self): - # Without height/width the pipeline defaults to 1024x1024 and squishes non-square - # input (the abba mixed-seam regression). Both already multiples of 16 -> unchanged. - from remove_ai_watermarks._internal.watermark_remover import _build_qwen_kwargs - - kwargs = _build_qwen_kwargs( - _StubImage(2816, 1536), strength=0.25, num_inference_steps=40, true_cfg_scale=4.0, generator=None - ) - assert kwargs["width"] == 2816 - assert kwargs["height"] == 1536 - - def test_qwen_target_size_floors_to_multiple_of_16(self): - from remove_ai_watermarks._internal.watermark_remover import _qwen_target_size - - assert _qwen_target_size(2816, 1536) == (2816, 1536) # already /16 - assert _qwen_target_size(1122, 1402) == (1120, 1392) # floored - assert _qwen_target_size(10, 10) == (16, 16) # min clamp, never 0 - - def test_qwen_model_id_is_qwen_image(self): - from remove_ai_watermarks._internal.watermark_profiles import QWEN_MODEL_ID - - assert QWEN_MODEL_ID == "Qwen/Qwen-Image" + assert calls["pipeline"].get("add_watermarker") is False + assert "add_watermarker" not in calls["controlnet"] class TestResolveStrength: - """resolve_strength applies the vendor default only when strength is unset.""" + """resolve_strength answers for sdxl-zimage and defers for qwen-zimage.""" - def test_none_is_vendor_adaptive(self): - # No vendor -> unknown default; OpenAI lower, Google == unknown. The sdxl/controlnet - # pipelines share this ladder (the certified controlnet floors); qwen has its own - # (see test_qwen_pipeline_uses_its_own_higher_ladder). - assert resolve_strength(None) == UNKNOWN_STRENGTH - assert resolve_strength(None, "openai") == OPENAI_STRENGTH - assert resolve_strength(None, "google") == GEMINI_STRENGTH - assert resolve_strength(None, None) == UNKNOWN_STRENGTH - # An unrecognized vendor string falls through to the unknown default. - assert resolve_strength(None, "adobe") == UNKNOWN_STRENGTH - # sdxl/controlnet pipelines (and the "default" alias) use the same shared ladder. - assert resolve_strength(None, "google", "controlnet") == GEMINI_STRENGTH - assert resolve_strength(None, "google", "sdxl") == GEMINI_STRENGTH + def test_qwen_zimage_answers_from_the_resolution_curve(self): + """The function is total: it owns both policies rather than returning None. - def test_qwen_pipeline_uses_its_own_higher_ladder(self): - # Qwen's certified Gemini floor (0.25) is HIGHER than controlnet's (0.15); OpenAI - # matches (0.10). Unknown vendor on qwen tracks the higher Gemini value. This retires - # the old manual "pass --strength 0.25 for Gemini on qwen" workaround. - from remove_ai_watermarks._internal.watermark_profiles import QWEN_GEMINI_STRENGTH, QWEN_OPENAI_STRENGTH + qwen-zimage picks strength from image area, so it takes the size. Returning + None for it would push that branch onto every caller and leave one of the two + strength policies living outside this module. The vendor is ignored here on + purpose - the curve, not the issuer, is what was calibrated. + """ + assert resolve_strength(None, "google", "qwen-zimage", size=(2000, 1850)) == pytest.approx(0.154) + assert resolve_strength(None, None, "qwen-zimage", size=(600, 500)) == pytest.approx(0.084) - assert QWEN_GEMINI_STRENGTH == 0.25 - assert QWEN_OPENAI_STRENGTH == 0.10 - assert resolve_strength(None, "google", "qwen") == QWEN_GEMINI_STRENGTH - assert resolve_strength(None, "openai", "qwen") == QWEN_OPENAI_STRENGTH - assert resolve_strength(None, None, "qwen") == QWEN_GEMINI_STRENGTH # unknown -> higher floor - assert resolve_strength(None, "google", "qwen") > resolve_strength(None, "google", "controlnet") - # An explicit strength still wins on qwen. - assert resolve_strength(0.12, "google", "qwen") == 0.12 + def test_qwen_zimage_without_a_size_fails_loudly(self): + """A missing size must not silently fall back to some vendor value.""" + with pytest.raises(ValueError, match="size is required"): + resolve_strength(None, "google", "qwen-zimage") - def test_ladder_is_the_certified_controlnet_floors(self): - # The unified ladder == the oracle-certified controlnet floors. Lowered on the - # 2026-06-14 Modal re-test (OpenAI 0.10, Google/unknown 0.15); Google is the - # more-robust watermark, so it is higher. - assert OPENAI_STRENGTH == 0.10 - assert GEMINI_STRENGTH == 0.15 - assert UNKNOWN_STRENGTH == 0.15 - assert OPENAI_STRENGTH < GEMINI_STRENGTH + def test_sdxl_zimage_uses_its_flat_vendor_ladder(self): - def test_default_strength_alias_is_unknown_vendor_value(self): - assert DEFAULT_STRENGTH == UNKNOWN_STRENGTH - assert OPENAI_STRENGTH < UNKNOWN_STRENGTH + assert SDXL_ZIMAGE_OPENAI_STRENGTH == 0.15 + assert SDXL_ZIMAGE_GEMINI_STRENGTH == 0.25 + assert SDXL_ZIMAGE_UNKNOWN_STRENGTH == SDXL_ZIMAGE_GEMINI_STRENGTH + assert resolve_strength(None, "openai", "sdxl-zimage") == SDXL_ZIMAGE_OPENAI_STRENGTH + assert resolve_strength(None, "google", "sdxl-zimage") == SDXL_ZIMAGE_GEMINI_STRENGTH + # An unrecognised issuer takes the stricter Gemini value, not the OpenAI one. + assert resolve_strength(None, "adobe", "sdxl-zimage") == SDXL_ZIMAGE_UNKNOWN_STRENGTH + assert resolve_strength(None, None, "sdxl-zimage") == SDXL_ZIMAGE_UNKNOWN_STRENGTH def test_strength_default_help_derives_from_constants(self): - # The CLI --strength help is built from this, so it can never drift from the ladder. + h = strength_default_help() - assert str(OPENAI_STRENGTH) in h - assert str(GEMINI_STRENGTH) in h - assert str(UNKNOWN_STRENGTH) in h + assert str(SDXL_ZIMAGE_OPENAI_STRENGTH) in h + assert str(SDXL_ZIMAGE_GEMINI_STRENGTH) in h def test_explicit_value_overrides_vendor(self): - assert resolve_strength(0.3) == 0.3 - assert resolve_strength(0.3, "openai") == 0.3 + + assert resolve_strength(0.3, "openai", "sdxl-zimage") == 0.3 + assert resolve_strength(0.3, None, "qwen-zimage") == 0.3 def test_explicit_zero_is_respected_not_treated_as_unset(self): - # 0.0 is falsy but explicit -- must not fall through to the vendor default + # 0.0 is falsy but explicit -- it must not fall through to the vendor default # (the old `strength or DEFAULT` bug would have). Range validation lives in # remove_watermark, not here. - assert resolve_strength(0.0) == 0.0 - assert resolve_strength(0.0, "google") == 0.0 + + assert resolve_strength(0.0, "google", "sdxl-zimage") == 0.0 + assert resolve_strength(0.0, None, "qwen-zimage") == 0.0 class TestVendorForStrength: @@ -467,54 +332,3 @@ class TestPlatformPaths: # If we get here without error, asset loading works assert engine._alpha_small.shape == (48, 48) assert engine._alpha_large.shape == (96, 96) - - -class TestFp16VaeFix: - """The plain SDXL img2img pipeline must swap in the fp16-fixed VAE on fp16 - GPUs to avoid the NaN/all-black decode (issue #29). Pure decision logic, no - torch or model download needed.""" - - DEFAULT = "stabilityai/stable-diffusion-xl-base-1.0" - - def test_default_sdxl_on_fp16_needs_fix(self): - from remove_ai_watermarks._internal.watermark_remover import _needs_fp16_vae_fix - - assert _needs_fp16_vae_fix(self.DEFAULT, self.DEFAULT, is_fp16=True) is True - - def test_fp32_does_not_need_fix(self): - """cpu/mps run fp32, where the stock SDXL VAE is fine.""" - from remove_ai_watermarks._internal.watermark_remover import _needs_fp16_vae_fix - - assert _needs_fp16_vae_fix(self.DEFAULT, self.DEFAULT, is_fp16=False) is False - - def test_non_default_model_keeps_own_vae(self): - """A custom (non-SDXL) checkpoint must not get the SDXL-specific VAE.""" - from remove_ai_watermarks._internal.watermark_remover import _needs_fp16_vae_fix - - assert _needs_fp16_vae_fix("runwayml/stable-diffusion-v1-5", self.DEFAULT, is_fp16=True) is False - - -class TestDegenerateOutputGuard: - """The fp16 black-output safety net (#29/#41): detect an all-black/NaN frame so - ``remove_watermark`` can retry in fp32. Pure image statistics, no model needed.""" - - def test_all_black_is_degenerate(self): - from remove_ai_watermarks._internal.watermark_remover import _is_degenerate_image - - black = Image.fromarray(np.zeros((64, 64, 3), np.uint8)) - assert _is_degenerate_image(black) is True - - def test_normal_image_is_not_degenerate(self): - from remove_ai_watermarks._internal.watermark_remover import _is_degenerate_image - - rng = np.random.default_rng(0) - normal = Image.fromarray(rng.integers(0, 256, (64, 64, 3), dtype=np.uint8)) - assert _is_degenerate_image(normal) is False - - def test_dark_but_textured_image_is_not_degenerate(self): - """A legitimately dark photo with real detail must NOT be flagged (variance guard).""" - from remove_ai_watermarks._internal.watermark_remover import _is_degenerate_image - - rng = np.random.default_rng(1) - dark = Image.fromarray(rng.integers(0, 40, (64, 64, 3), dtype=np.uint8)) - assert _is_degenerate_image(dark) is False diff --git a/tests/test_qwen_zimage_pipeline.py b/tests/test_qwen_zimage_pipeline.py index 366a42d..c8d4ae8 100644 --- a/tests/test_qwen_zimage_pipeline.py +++ b/tests/test_qwen_zimage_pipeline.py @@ -530,12 +530,10 @@ def test_profile_defaults_to_four_global_steps(): ) assert normalize_profile("qwen-zimage") == "qwen-zimage" - assert resolve_steps(None, "qwen-zimage") == 4 - assert resolve_steps(None, "controlnet") == 50 - assert resolve_steps(12, "qwen-zimage") == 12 - assert resolve_seed(None, "qwen-zimage") == 0 - assert resolve_seed(None, "controlnet") is None - assert resolve_seed(17, "qwen-zimage") == 17 + assert resolve_steps(None) == 4 + assert resolve_steps(12) == 12 + assert resolve_seed(None) == 0 + assert resolve_seed(17) == 17 def test_cli_exposes_qwen_zimage_profile(): @@ -586,7 +584,7 @@ def test_watermark_remover_dispatches_to_full_pipeline(tmp_path, monkeypatch): runtime = MagicMock() runtime.run.return_value = Image.new("RGB", (64, 48), (50, 60, 70)) - remover = WatermarkRemover(device="cpu", pipeline="qwen-zimage") + remover = WatermarkRemover(device="cuda", pipeline="qwen-zimage") monkeypatch.setattr(remover, "_load_qwen_zimage_pipeline", lambda: runtime) assert remover.model_id == "Qwen/Qwen-Image-2512 + Tongyi-MAI/Z-Image-Turbo" @@ -612,7 +610,7 @@ def test_watermark_remover_dispatches_qwen_tiling_to_full_pipeline(tmp_path, mon runtime = MagicMock() runtime.run.return_value = Image.new("RGB", (96, 80), (50, 60, 70)) - remover = WatermarkRemover(device="cpu", pipeline="qwen-zimage") + remover = WatermarkRemover(device="cuda", pipeline="qwen-zimage") monkeypatch.setattr(remover, "_load_qwen_zimage_pipeline", lambda: runtime) remover.remove_watermark( @@ -746,11 +744,11 @@ def test_qwen_zimage_rejects_runtime_knobs_that_change_fixed_graph(tmp_path, mon _mock_watermark_runtime_deps(monkeypatch) with pytest.raises(ValueError, match="fixed Qwen-Image-2512"): - WatermarkRemover(model_id="custom/model", device="cpu", pipeline="qwen-zimage") + WatermarkRemover(model_id="custom/model", device="cuda", pipeline="qwen-zimage") source = tmp_path / "source.png" Image.new("RGB", (64, 48)).save(source) - remover = WatermarkRemover(device="cpu", pipeline="qwen-zimage") + remover = WatermarkRemover(device="cuda", pipeline="qwen-zimage") with pytest.raises(ValueError, match=r"CFG 1\.0"): remover.remove_watermark(source, guidance_scale=2.0) with pytest.raises(ValueError, match="requires 4 steps"): @@ -786,10 +784,11 @@ def test_sdxl_zimage_strength_is_vendor_adaptive_and_leaves_other_profiles_alone 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. + # An explicit value still wins, and qwen-zimage is untouched by this ladder: it + # defers to its resolution curve rather than to a vendor value. 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) + assert resolve_strength(None, "openai", "qwen-zimage", size=(2000, 1850)) == pytest.approx(0.154) + assert resolve_strength(None, "google", "qwen-zimage", size=(2000, 1850)) == pytest.approx(0.154) def test_sdxl_zimage_shares_the_four_step_seed_and_step_contract(): @@ -800,10 +799,8 @@ def test_sdxl_zimage_shares_the_four_step_seed_and_step_contract(): ) 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 + assert resolve_steps(None) == 4 + assert resolve_seed(None) == 0 def test_sdxl_requested_steps_compensate_for_the_diffusers_truncation(): diff --git a/tests/test_watermark_profiles.py b/tests/test_watermark_profiles.py deleted file mode 100644 index 85964f1..0000000 --- a/tests/test_watermark_profiles.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Pure tests for the strength/steps profile helpers (no model, no torch needed).""" - -from __future__ import annotations - -import pytest - -from remove_ai_watermarks._internal.watermark_profiles import resolve_strength, viable_steps - - -class TestViableSteps: - """Guards the crash found by the release smoke matrix on 2026-07-19. - - diffusers derives its img2img timesteps as ``int(steps * strength)``. When that - rounds to zero the pipeline builds an empty tensor and dies deep inside attention - with "cannot reshape tensor of 0 elements into shape [0, -1, 1, 512]". At the - default strength 0.15 that was every ``--steps`` below 7, reachable with entirely - valid CLI arguments and no special flags. - """ - - @pytest.mark.parametrize( - ("steps", "strength"), - [(1, 0.15), (2, 0.15), (5, 0.15), (6, 0.15), (5, 0.10), (9, 0.10), (1, 0.5)], - ) - def test_never_returns_a_count_that_denoises_zero_steps(self, steps: int, strength: float): - assert int(viable_steps(steps, strength) * strength) >= 1 - - @pytest.mark.parametrize( - ("steps", "strength"), - [(50, 0.15), (20, 0.15), (7, 0.15), (10, 0.10), (2, 0.5), (50, 1.0)], - ) - def test_leaves_a_workable_count_untouched(self, steps: int, strength: float): - assert viable_steps(steps, strength) == steps - - def test_raises_only_to_the_minimum_needed(self): - # strength 0.15 needs 7 (int(7*0.15)==1); it must not jump to some larger default. - assert viable_steps(5, 0.15) == 7 - assert viable_steps(1, 0.10) == 10 - - def test_the_vendor_defaults_all_have_a_reachable_floor(self): - for vendor in (None, "openai", "google"): - strength = resolve_strength(None, vendor) - assert int(viable_steps(1, strength) * strength) >= 1 - - @pytest.mark.parametrize("strength", [0.0, -0.1]) - def test_a_non_positive_strength_cannot_loop_or_divide_by_zero(self, strength: float): - # No denoising is possible at all here; return the caller's value rather than - # dividing by zero or spinning. - assert viable_steps(20, strength) == 20 diff --git a/uv.lock b/uv.lock index e93104e..af7caab 100644 --- a/uv.lock +++ b/uv.lock @@ -3331,7 +3331,7 @@ wheels = [ [[package]] name = "remove-ai-watermarks" -version = "0.23.0" +version = "0.24.0" source = { editable = "." } dependencies = [ { name = "c2pa-python" },