CalvinBackup/remove-ai-watermarks
1
0
Fork 0
mirror of https://github.com/wiltodelta/remove-ai-watermarks.git synced 2026-08-10 16:10:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

Keep only the two-stage profiles and make CUDA a precondition

Browse Source 
qwen-zimage becomes the default and sdxl-zimage the only alternative. The
controlnet, sdxl, qwen and default profiles are gone, and with them the CPU and
MPS paths for invisible-watermark removal: neither matched the two-stage
recipe's face preservation, so keeping them advertised a quality this library no
longer delivers. Visible-mark removal and every identify command still run
anywhere.

Retired names are rejected rather than remapped. Silently routing --pipeline
sdxl onward would run an old script at a different strength, on a different
model, at a different quality, and report success.

CUDA is now checked when the remover is constructed instead of when the model
loads. Auto-detection cheerfully returned mps on a Mac, so the failure arrived
several layers down, after the dependency check and the pipeline import, in a
message naming whichever internal pipeline happened to raise. _DEVICES collapses
to {"cuda"} and the cpu/mps float32 branch goes with it.

resolve_strength stays total. It briefly returned None for qwen-zimage, meaning
"ask the resolution curve", which pushed a branch onto both callers and left one
of the two strength policies outside the strength module; the CLI copy had
already grown an `or 0.0` guarding a path its own comment called unreachable. It
now takes the image size and answers for both profiles, so the displayed value
cannot drift from the executed one.

Deletion fallout removed with it: img2img_runner and progress.py (the MPS
recovery path and its progress monitor had no callers left), viable_steps, the
fp16 degenerate-output retry, the fp16 VAE fix, and the Qwen img2img call
builders. try_empty_device_cache moved into watermark_remover rather than
leaving a module whose docstring outlived its code. _HAS_DIFFUSERS routes
through optional_deps.module_available, which is what the rest of the library
uses and what correctly rejects a pruned namespace remnant.

--steps, --guidance-scale and --model now have exactly one legal value each and
are still accepted at parse time, then rejected in remove(). Their help text
says so, but validating them beside the option would be better.

Not addressed, and worth its own decision: invisible_engine forces
min_resolution to 0 for both profiles, so the --min-resolution floor, --upscaler,
_esrgan_upscale, upscaler.py and the esrgan extra are all unreachable.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
2026-08-03 12:16:14 -07:00
co-authored by Claude Opus 5
parent 3d43bac6a5
commit b0ca2054f6
23 changed files with 331 additions and 1458 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/remove_ai_watermarks/__init__.py
+1 -1
View File
@@ -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__",
src/remove_ai_watermarks/_internal/img2img_runner.py
-136
View File
@@ -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()
src/remove_ai_watermarks/_internal/progress.py
-212
View File
@@ -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()
src/remove_ai_watermarks/_internal/qwen_zimage_pipeline.py
+1 -1
View File
@@ -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)
)
src/remove_ai_watermarks/_internal/watermark_profiles.py
+48 -52
View File
@@ -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:
src/remove_ai_watermarks/_internal/watermark_remover.py
+53 -359
View File
@@ -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:
src/remove_ai_watermarks/cli.py
+50 -65
View File
@@ -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,
src/remove_ai_watermarks/invisible_engine.py
+13 -17
View File
@@ -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