mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-07 06:28:36 +02:00
Sync dependency updates with current main
This commit is contained in:
@@ -25,7 +25,7 @@ _os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
|
||||
_warnings.filterwarnings("ignore", message=r".*ImageProcessorFast.*")
|
||||
|
||||
|
||||
__version__ = "0.19.0"
|
||||
__version__ = "0.20.0"
|
||||
|
||||
__all__ = ["__version__", "remove_visible", "visible_provenance"]
|
||||
|
||||
|
||||
@@ -21,6 +21,8 @@ import click
|
||||
from remove_ai_watermarks import __version__, image_io, watermark_registry
|
||||
from remove_ai_watermarks.noai.constants import SUPPORTED_FORMATS
|
||||
from remove_ai_watermarks.noai.watermark_profiles import (
|
||||
resolve_seed,
|
||||
resolve_steps,
|
||||
resolve_strength,
|
||||
strength_default_help,
|
||||
vendor_for_strength,
|
||||
@@ -137,6 +139,23 @@ def _validate_image(path: Path) -> Path:
|
||||
return path
|
||||
|
||||
|
||||
def _resolved_strength_for_display(
|
||||
source: Path,
|
||||
strength: float | None,
|
||||
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
|
||||
|
||||
from remove_ai_watermarks.noai.qwen_zimage_pipeline import resolution_adaptive_denoise
|
||||
|
||||
with Image.open(source) as image:
|
||||
return resolution_adaptive_denoise(image.width, image.height)
|
||||
return resolve_strength(strength, vendor, pipeline)
|
||||
|
||||
|
||||
# Shared option decorator for commands that run the invisible-watermark pipeline.
|
||||
# Both cmd_invisible and cmd_all expose this flag; defining it once avoids
|
||||
# copy-paste drift.
|
||||
@@ -184,8 +203,9 @@ _adaptive_polish_option = click.option(
|
||||
default=True,
|
||||
help="Restore the input's detail level after removal (capped unsharp + edge-masked grain "
|
||||
"targeting the input's sharpness, sparing text), countering the over-smoothed look. ON by "
|
||||
"default; it self-limits where there is no detail deficit (text/flat graphics), so it is a "
|
||||
"no-op there. Pass --no-adaptive-polish to disable. Independent of --unsharp/--humanize.",
|
||||
"default except for qwen-zimage, whose upstream-matching output is left unchanged; it "
|
||||
"self-limits where there is no detail deficit (text/flat graphics). Pass --adaptive-polish "
|
||||
"or --no-adaptive-polish to override. Independent of --unsharp/--humanize.",
|
||||
)
|
||||
|
||||
|
||||
@@ -204,7 +224,7 @@ def _tile_options(f: Any) -> Any:
|
||||
"--tile-size",
|
||||
type=int,
|
||||
default=1024,
|
||||
help="Tile dimension in px for --tile (SDXL's training size). Default 1024.",
|
||||
help="Tile dimension in px for --tile. Default 1024.",
|
||||
)(f)
|
||||
return click.option(
|
||||
"--tile/--no-tile",
|
||||
@@ -227,8 +247,8 @@ _guidance_scale_option = click.option(
|
||||
"--guidance-scale",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Classifier-free guidance scale (CFG). Default: 7.5 (the library default). "
|
||||
"Lower = follow the prompt less / stay closer to the input.",
|
||||
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.",
|
||||
)
|
||||
|
||||
|
||||
@@ -255,13 +275,15 @@ def _normalize_pipeline(ctx: click.Context, param: click.Parameter, value: str |
|
||||
# ``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", "default"]
|
||||
_PIPELINE_CHOICES = ["sdxl", "controlnet", "qwen", "qwen-zimage", "default"]
|
||||
_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 (does not fit MPS). ('default' is an OUTDATED alias for 'sdxl'.)"
|
||||
"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'.)"
|
||||
)
|
||||
|
||||
# Shared --pipeline / --strength decorators so the three diffusion commands
|
||||
@@ -336,6 +358,19 @@ def _resolve_auto_polish(auto: bool, adaptive_polish: bool) -> bool:
|
||||
return adaptive_polish
|
||||
|
||||
|
||||
def _resolve_profile_polish(auto: bool, adaptive_polish: bool, pipeline: str) -> bool:
|
||||
"""Keep the upstream qwen-zimage output unchanged unless polish was explicit."""
|
||||
adaptive_polish = _resolve_auto_polish(auto, adaptive_polish)
|
||||
if pipeline != "qwen-zimage":
|
||||
return adaptive_polish
|
||||
ctx = click.get_current_context(silent=True)
|
||||
if ctx is None:
|
||||
return adaptive_polish
|
||||
if ctx.get_parameter_source("adaptive_polish") == click.core.ParameterSource.DEFAULT:
|
||||
return False
|
||||
return adaptive_polish
|
||||
|
||||
|
||||
def _warn_if_esrgan_unavailable(upscaler: str) -> None:
|
||||
"""Tell the user once if ``--upscaler esrgan`` will silently fall back to Lanczos.
|
||||
|
||||
@@ -827,7 +862,12 @@ def cmd_erase(
|
||||
"-o", "--output", type=click.Path(path_type=Path), default=None, help="Output path (default: <source>_clean.<ext>)."
|
||||
)
|
||||
@_strength_option
|
||||
@click.option("--steps", type=int, default=50, help="Number of denoising steps. Default: 50.")
|
||||
@click.option(
|
||||
"--steps",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Number of denoising steps. Default: 4 for qwen-zimage, 50 otherwise.",
|
||||
)
|
||||
@_pipeline_option
|
||||
@click.option(
|
||||
"--device",
|
||||
@@ -835,7 +875,12 @@ def cmd_erase(
|
||||
default="auto",
|
||||
help="Inference device.",
|
||||
)
|
||||
@click.option("--seed", type=int, default=None, help="Random seed for reproducibility.")
|
||||
@click.option(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Random seed for reproducibility. Default: 0 for qwen-zimage, random otherwise.",
|
||||
)
|
||||
@click.option("--hf-token", type=str, default=None, help="HuggingFace API token.")
|
||||
@click.option(
|
||||
"--humanize", type=float, default=0.0, help="Analog Humanizer film grain intensity (0 = off, typical: 2.0-6.0)."
|
||||
@@ -862,7 +907,7 @@ def cmd_invisible(
|
||||
source: Path,
|
||||
output: Path | None,
|
||||
strength: float | None,
|
||||
steps: int,
|
||||
steps: int | None,
|
||||
pipeline: str,
|
||||
device: str,
|
||||
seed: int | None,
|
||||
@@ -898,8 +943,10 @@ 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)
|
||||
_warn_if_esrgan_unavailable(upscaler)
|
||||
adaptive_polish = _resolve_auto_polish(auto, adaptive_polish)
|
||||
adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline)
|
||||
if output is None:
|
||||
output = source.with_stem(source.stem + "_clean")
|
||||
|
||||
@@ -929,7 +976,7 @@ def cmd_invisible(
|
||||
vendor = vendor_for_strength(source)
|
||||
console.print(f" Input: {source.name}")
|
||||
console.print(f" Pipeline: {pipeline}")
|
||||
console.print(f" Strength: {resolve_strength(strength, vendor, pipeline)} Steps: {steps}")
|
||||
console.print(f" Strength: {_resolved_strength_for_display(source, strength, vendor, pipeline)} Steps: {steps}")
|
||||
|
||||
t0 = time.monotonic()
|
||||
result_path = engine.remove_watermark(
|
||||
@@ -1101,7 +1148,12 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo
|
||||
@_visible_backend_option
|
||||
@_visible_sensitivity_option
|
||||
@_strength_option
|
||||
@click.option("--steps", type=int, default=50, help="Number of denoising steps for invisible removal.")
|
||||
@click.option(
|
||||
"--steps",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Number of denoising steps. Default: 4 for qwen-zimage, 50 otherwise.",
|
||||
)
|
||||
@_pipeline_option
|
||||
@_model_option
|
||||
@click.option(
|
||||
@@ -1110,7 +1162,12 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo
|
||||
default="auto",
|
||||
help="Inference device.",
|
||||
)
|
||||
@click.option("--seed", type=int, default=None, help="Random seed for reproducibility.")
|
||||
@click.option(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Random seed for reproducibility. Default: 0 for qwen-zimage, random otherwise.",
|
||||
)
|
||||
@click.option("--hf-token", type=str, default=None, help="HuggingFace API token.")
|
||||
@click.option(
|
||||
"--humanize", type=float, default=0.0, help="Analog Humanizer film grain intensity (0 = off, typical: 2.0-6.0)."
|
||||
@@ -1138,7 +1195,7 @@ def cmd_all(
|
||||
backend: str,
|
||||
sensitivity: str,
|
||||
strength: float | None,
|
||||
steps: int,
|
||||
steps: int | None,
|
||||
pipeline: str,
|
||||
model: str | None,
|
||||
device: str,
|
||||
@@ -1169,8 +1226,10 @@ def cmd_all(
|
||||
"""
|
||||
_banner()
|
||||
source = _validate_image(source)
|
||||
steps = resolve_steps(steps, pipeline)
|
||||
seed = resolve_seed(seed, pipeline)
|
||||
_warn_if_esrgan_unavailable(upscaler)
|
||||
adaptive_polish = _resolve_auto_polish(auto, adaptive_polish)
|
||||
adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline)
|
||||
|
||||
if output is None:
|
||||
output = source.with_stem(source.stem + "_clean")
|
||||
@@ -1260,7 +1319,9 @@ def cmd_all(
|
||||
# already lost its C2PA to the visible-removal pass, so reading it would
|
||||
# always resolve to the unknown-vendor default.
|
||||
vendor = vendor_for_strength(source)
|
||||
console.print(f" Strength: {resolve_strength(strength, vendor, pipeline)} Steps: {steps}")
|
||||
console.print(
|
||||
f" Strength: {_resolved_strength_for_display(source, strength, vendor, pipeline)} Steps: {steps}"
|
||||
)
|
||||
inv_engine.remove_watermark(
|
||||
image_path=tmp_path,
|
||||
output_path=tmp_path,
|
||||
@@ -1521,7 +1582,12 @@ def _process_batch_image(
|
||||
"--mode", type=click.Choice(["visible", "invisible", "metadata", "all"]), default="visible", help="Processing mode."
|
||||
)
|
||||
@_strength_option
|
||||
@click.option("--steps", type=int, default=50, help="Number of denoising steps (invisible mode).")
|
||||
@click.option(
|
||||
"--steps",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Number of denoising steps. Default: 4 for qwen-zimage, 50 otherwise.",
|
||||
)
|
||||
@_visible_backend_option
|
||||
@_visible_sensitivity_option
|
||||
@click.option(
|
||||
@@ -1534,7 +1600,12 @@ def _process_batch_image(
|
||||
default="auto",
|
||||
help="Inference device.",
|
||||
)
|
||||
@click.option("--seed", type=int, default=None, help="Random seed for reproducibility.")
|
||||
@click.option(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Random seed for reproducibility. Default: 0 for qwen-zimage, random otherwise.",
|
||||
)
|
||||
@click.option("--hf-token", type=str, default=None, help="HuggingFace API token.")
|
||||
@click.option(
|
||||
"--max-resolution",
|
||||
@@ -1559,7 +1630,7 @@ def cmd_batch(
|
||||
mode: str,
|
||||
output_dir: Path | None,
|
||||
strength: float | None,
|
||||
steps: int,
|
||||
steps: int | None,
|
||||
pipeline: str,
|
||||
device: str,
|
||||
seed: int | None,
|
||||
@@ -1599,7 +1670,9 @@ def cmd_batch(
|
||||
console.print(f" Mode: {mode}")
|
||||
if mode in ("invisible", "all"):
|
||||
_warn_if_esrgan_unavailable(upscaler)
|
||||
adaptive_polish = _resolve_auto_polish(auto, adaptive_polish)
|
||||
adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline)
|
||||
steps = resolve_steps(steps, pipeline)
|
||||
seed = resolve_seed(seed, pipeline)
|
||||
options = _BatchOptions(
|
||||
strength=strength,
|
||||
steps=steps,
|
||||
|
||||
@@ -19,7 +19,12 @@ import warnings
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from .noai.watermark_profiles import DEFAULT_MODEL_ID as DEFAULT_SDXL_MODEL_ID
|
||||
from .noai.watermark_profiles import (
|
||||
DEFAULT_MODEL_ID as DEFAULT_SDXL_MODEL_ID,
|
||||
)
|
||||
from .noai.watermark_profiles import (
|
||||
resolve_seed,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
@@ -105,7 +110,9 @@ class InvisibleEngine:
|
||||
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). "default" aliases "sdxl".
|
||||
structure preservation but CUDA/cloud-class), or "qwen-zimage"
|
||||
(Qwen-Image-2512 Lightning + Canny, then SAM-masked Z-Image face
|
||||
repair; CUDA-only). "default" aliases "sdxl".
|
||||
hf_token: HuggingFace API token.
|
||||
progress_callback: Optional callback for progress messages.
|
||||
controlnet_conditioning_scale: ControlNet structure-preservation
|
||||
@@ -161,7 +168,7 @@ class InvisibleEngine:
|
||||
image_path: Path,
|
||||
output_path: Path | None = None,
|
||||
strength: float | None = None,
|
||||
num_inference_steps: int = 100,
|
||||
num_inference_steps: int | None = None,
|
||||
guidance_scale: float | None = None,
|
||||
seed: int | None = None,
|
||||
humanize: float = 0.0,
|
||||
@@ -182,9 +189,12 @@ class InvisibleEngine:
|
||||
output_path: Output path (None = overwrite source).
|
||||
strength: Denoising strength (0.0-1.0). None -> the vendor-adaptive
|
||||
default.
|
||||
num_inference_steps: Number of denoising steps.
|
||||
num_inference_steps: Number of denoising steps. None keeps the existing
|
||||
100-step library default, except qwen-zimage uses its required
|
||||
four-step Lightning schedule.
|
||||
guidance_scale: Classifier-free guidance scale.
|
||||
seed: Random seed for reproducibility.
|
||||
seed: Random seed for reproducibility. None resolves to 0 for
|
||||
qwen-zimage and stays random for the other profiles.
|
||||
humanize: Intensity of Analog Humanizer film grain (0 = off).
|
||||
unsharp: Final unsharp-mask sharpening strength (0 = off, default).
|
||||
Applied last to counter the soft / over-smoothed look of the
|
||||
@@ -222,6 +232,12 @@ class InvisibleEngine:
|
||||
"""
|
||||
import tempfile
|
||||
|
||||
if num_inference_steps is None:
|
||||
profile = getattr(self._remover, "model_profile", None)
|
||||
num_inference_steps = 4 if profile == "qwen-zimage" else 100
|
||||
profile = getattr(self._remover, "model_profile", "controlnet")
|
||||
seed = resolve_seed(seed, profile)
|
||||
|
||||
from PIL import Image, ImageOps
|
||||
|
||||
# Resolution policy: a max_resolution cap (0 = none) bounds memory on huge
|
||||
@@ -242,7 +258,18 @@ class InvisibleEngine:
|
||||
# reassigned to the resized copy below; PIL resize returns a new object).
|
||||
reference_pil = image
|
||||
|
||||
target = _target_size(image.width, image.height, max_resolution, min_resolution)
|
||||
# qwen-zimage operates at the input's native geometry in its reference graph.
|
||||
# Keep an explicit max cap available for callers, but do not apply the SDXL
|
||||
# 1024px minimum-resolution floor to this profile.
|
||||
effective_min_resolution = (
|
||||
0 if getattr(self._remover, "model_profile", None) == "qwen-zimage" else min_resolution
|
||||
)
|
||||
target = _target_size(
|
||||
image.width,
|
||||
image.height,
|
||||
max_resolution,
|
||||
effective_min_resolution,
|
||||
)
|
||||
if target is not None:
|
||||
upscaling = max(target) > max(image.width, image.height)
|
||||
if self._progress_callback:
|
||||
@@ -356,9 +383,12 @@ class InvisibleEngine:
|
||||
input_dir: Path,
|
||||
output_dir: Path,
|
||||
strength: float | None = None,
|
||||
steps: int = 50,
|
||||
steps: int | None = None,
|
||||
) -> list[Path]:
|
||||
"""Remove invisible watermarks from all images in a directory."""
|
||||
if steps is None:
|
||||
profile = getattr(self._remover, "model_profile", None)
|
||||
steps = 4 if profile == "qwen-zimage" else 50
|
||||
return self._remover.remove_watermark_batch(
|
||||
input_dir=input_dir,
|
||||
output_dir=output_dir,
|
||||
|
||||
@@ -0,0 +1,903 @@
|
||||
"""Qwen 2512 Canny regeneration followed by masked Z-Image face repair.
|
||||
|
||||
This profile ports the two-stage architecture used by cebeuq/Synthid-Bypass:
|
||||
|
||||
1. Qwen-Image-2512 img2img with the 4-step Lightning LoRA and the DiffSynth
|
||||
blockwise Canny ControlNet regenerates the whole image, optionally as
|
||||
overlapping feather-blended tiles for large inputs.
|
||||
2. Faces are detected on the original image, refined to masks with SAM, regenerated
|
||||
from the original face crops with Z-Image Turbo, and feathered into stage 1.
|
||||
|
||||
The runtime intentionally uses permissively licensed YuNet instead of the reference
|
||||
workflow's Ultralytics detector. All diffusion and segmentation models remain the same
|
||||
model families and the denoise formulas are direct ports of the reference custom node.
|
||||
"""
|
||||
|
||||
# DiffSynth, torch, transformers, and cv2 expose mostly untyped tensor/array APIs.
|
||||
# Keep the relaxation local to this optional ML boundary.
|
||||
# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportUnnecessaryComparison=false
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import hashlib
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import tempfile
|
||||
import urllib.request
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from remove_ai_watermarks.noai.watermark_profiles import resolve_seed
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
QWEN_IMAGE_2512_MODEL_ID = "Qwen/Qwen-Image-2512"
|
||||
QWEN_CANNY_CONTROLNET_MODEL_ID = "DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Canny"
|
||||
QWEN_LIGHTNING_MODEL_ID = "lightx2v/Qwen-Image-2512-Lightning"
|
||||
QWEN_LIGHTNING_PATTERN = "Qwen-Image-2512-Lightning-4steps-V1.0-bf16.safetensors"
|
||||
ZIMAGE_TURBO_MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
|
||||
SAM_MODEL_ID = "facebook/sam-vit-base"
|
||||
|
||||
YUNET_MODEL_URL = (
|
||||
"https://media.githubusercontent.com/media/opencv/opencv_zoo/main/"
|
||||
"models/face_detection_yunet/face_detection_yunet_2023mar.onnx"
|
||||
)
|
||||
YUNET_MODEL_NAME = "face_detection_yunet_2023mar.onnx"
|
||||
YUNET_MODEL_SHA256 = "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4"
|
||||
# The upstream graph's 0.2 threshold belongs to YOLO and does not transfer to
|
||||
# YuNet's score calibration. At 0.2 YuNet admitted background and decorative
|
||||
# false positives, multiplying the serial Z-Image face-stage cost. A 0.5 gate
|
||||
# retained all visible faces in the public and upstream comparison fixtures.
|
||||
YUNET_SCORE_THRESHOLD = 0.5
|
||||
|
||||
GLOBAL_STEPS = 4
|
||||
FACE_STEPS = 8
|
||||
GLOBAL_CFG = 1.0
|
||||
FACE_CFG = 1.0
|
||||
GLOBAL_CONTROLNET_SCALE = 1.0
|
||||
RESIDENT_FACE_MODEL_MIN_VRAM_GIB = 64.0
|
||||
|
||||
# The source graph uses normalized Canny thresholds 0.05 and 0.25. OpenCV takes
|
||||
# byte thresholds, so round 255*x to the matching integer values.
|
||||
_CANNY_LOW = 13
|
||||
_CANNY_HIGH = 64
|
||||
|
||||
# These strings intentionally preserve the reference workflow spelling. They are
|
||||
# model inputs, not user-facing copy, and changing them would change the port.
|
||||
_GLOBAL_PROMPT = "ultra clear and smoothe skin, spotless skin"
|
||||
_GLOBAL_NEGATIVE = "moles, freckes, high detail skin"
|
||||
_FACE_PROMPT = ""
|
||||
_FACE_NEGATIVE = "blurry, ugly, bad quality,"
|
||||
|
||||
|
||||
def resolve_face_model_residency(
|
||||
requested: bool | None,
|
||||
*,
|
||||
total_memory_gib: float,
|
||||
) -> bool:
|
||||
"""Keep the Z-Image stack resident when explicitly requested or safely sized."""
|
||||
if requested is not None:
|
||||
return requested
|
||||
return total_memory_gib >= RESIDENT_FACE_MODEL_MIN_VRAM_GIB
|
||||
|
||||
|
||||
def _pin_vram_managed_models(pipe: Any) -> None:
|
||||
"""Move the managed Z-Image stack to CUDA once and make offload a no-op."""
|
||||
model_names = ["text_encoder", "dit", "vae_encoder", "vae_decoder"]
|
||||
for name in model_names:
|
||||
model = getattr(pipe, name, None)
|
||||
if model is None:
|
||||
continue
|
||||
for module in model.modules():
|
||||
if not all(
|
||||
hasattr(module, attribute)
|
||||
for attribute in (
|
||||
"offload_dtype",
|
||||
"offload_device",
|
||||
"onload_dtype",
|
||||
"onload_device",
|
||||
"preparing_dtype",
|
||||
"preparing_device",
|
||||
"computation_dtype",
|
||||
"computation_device",
|
||||
)
|
||||
):
|
||||
continue
|
||||
module.offload_dtype = module.computation_dtype
|
||||
module.offload_device = module.computation_device
|
||||
module.onload_dtype = module.computation_dtype
|
||||
module.onload_device = module.computation_device
|
||||
module.preparing_dtype = module.computation_dtype
|
||||
module.preparing_device = module.computation_device
|
||||
pipe.load_models_to_device(model_names)
|
||||
|
||||
|
||||
def _cached_prompt_process(original_process: Any) -> Any:
|
||||
cache: dict[str, dict[str, Any]] = {}
|
||||
|
||||
def cached_process(
|
||||
runtime_pipe: Any,
|
||||
prompt: str,
|
||||
edit_image: Any = None,
|
||||
) -> dict[str, Any]:
|
||||
if edit_image is not None:
|
||||
return original_process(runtime_pipe, prompt, edit_image=edit_image)
|
||||
if prompt not in cache:
|
||||
cache[prompt] = original_process(runtime_pipe, prompt, edit_image=None)
|
||||
return cache[prompt]
|
||||
|
||||
return cached_process
|
||||
|
||||
|
||||
def _cache_static_prompt_embeddings(
|
||||
pipe: Any,
|
||||
output_params: tuple[str, ...],
|
||||
) -> bool:
|
||||
"""Memoize a prompt unit when its embedding depends only on static text."""
|
||||
for unit in pipe.units:
|
||||
if tuple(getattr(unit, "output_params", ())) == output_params:
|
||||
unit.process = _cached_prompt_process(unit.process)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _clamp(value: float, minimum: float, maximum: float) -> float:
|
||||
if maximum < minimum:
|
||||
minimum, maximum = maximum, minimum
|
||||
return max(minimum, min(maximum, value))
|
||||
|
||||
|
||||
def resolution_adaptive_denoise(
|
||||
width: int,
|
||||
height: int,
|
||||
*,
|
||||
adaptive_level: int = 6,
|
||||
denoise_min: float = 0.08,
|
||||
denoise_max: float = 0.15,
|
||||
) -> float:
|
||||
"""Port the reference resolution-based adaptive denoise calculation.
|
||||
|
||||
At neutral level 5, 0.30 MP maps to ``denoise_min`` and 3.70 MP maps to
|
||||
``denoise_max``. Levels above or below 5 add the same asymmetric spread as the
|
||||
reference custom node.
|
||||
"""
|
||||
image_mp = max(1.0, float(width) * float(height)) / 1_000_000.0
|
||||
normalized = _clamp((image_mp - 0.30) / (3.70 - 0.30), 0.0, 1.0)
|
||||
|
||||
minimum = float(denoise_min)
|
||||
maximum = float(denoise_max)
|
||||
if maximum < minimum:
|
||||
minimum, maximum = maximum, minimum
|
||||
denoise_range = maximum - minimum
|
||||
base = minimum + denoise_range * normalized
|
||||
|
||||
level = int(adaptive_level)
|
||||
if level >= 5:
|
||||
offset = ((float(level) - 5.0) / 5.0) * denoise_range * 0.285714
|
||||
else:
|
||||
offset = -((5.0 - float(level)) / 4.0) * denoise_range * 0.257143
|
||||
return _clamp(base + offset, 0.0001, 1.0)
|
||||
|
||||
|
||||
def largest_face_denoise(
|
||||
boxes: list[tuple[int, int, int, int]],
|
||||
image_size: tuple[int, int],
|
||||
*,
|
||||
base_denoise: float = 0.10,
|
||||
adaptive_ratio: float = 0.03,
|
||||
denoise_min: float = 0.05,
|
||||
denoise_max: float = 0.28,
|
||||
) -> float:
|
||||
"""Scale face denoise from the largest face area, matching reference mode."""
|
||||
width, height = image_size
|
||||
image_area = max(1.0, float(width) * float(height))
|
||||
largest_ratio = 0.0
|
||||
for x1, y1, x2, y2 in boxes:
|
||||
box_area = max(0.0, float(x2 - x1)) * max(0.0, float(y2 - y1))
|
||||
largest_ratio = max(largest_ratio, box_area / image_area)
|
||||
if largest_ratio <= 0.0:
|
||||
return _clamp(base_denoise, denoise_min, denoise_max)
|
||||
scaled = float(base_denoise) * (largest_ratio / max(1e-6, float(adaptive_ratio)))
|
||||
return _clamp(scaled, denoise_min, denoise_max)
|
||||
|
||||
|
||||
def _target_size(width: int, height: int) -> tuple[int, int]:
|
||||
"""Floor image dimensions to the /16 latent grid without changing aspect."""
|
||||
return max(16, (width // 16) * 16), max(16, (height // 16) * 16)
|
||||
|
||||
|
||||
def _resize_to_target(image: Image.Image) -> Image.Image:
|
||||
"""Resize pixels to the exact latent-grid dimensions passed to DiffSynth."""
|
||||
target = _target_size(image.width, image.height)
|
||||
if image.size == target:
|
||||
return image
|
||||
return image.resize(target, Image.Resampling.LANCZOS)
|
||||
|
||||
|
||||
def build_canny_control_image(image: Image.Image) -> Image.Image:
|
||||
"""Build the three-channel Canny conditioning image used by stage 1."""
|
||||
import cv2
|
||||
|
||||
rgb = np.asarray(image.convert("RGB"))
|
||||
gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
|
||||
edges = cv2.Canny(gray, _CANNY_LOW, _CANNY_HIGH)
|
||||
return Image.fromarray(np.repeat(edges[:, :, None], 3, axis=2))
|
||||
|
||||
|
||||
def build_global_kwargs(
|
||||
image: Image.Image,
|
||||
*,
|
||||
strength: float,
|
||||
seed: int | None,
|
||||
controlnet_input: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""Build the DiffSynth Qwen call shape without importing the ML runtime."""
|
||||
input_image = _resize_to_target(image)
|
||||
width, height = input_image.size
|
||||
return {
|
||||
"prompt": _GLOBAL_PROMPT,
|
||||
"negative_prompt": _GLOBAL_NEGATIVE,
|
||||
"cfg_scale": GLOBAL_CFG,
|
||||
"input_image": input_image,
|
||||
"denoising_strength": float(strength),
|
||||
"height": height,
|
||||
"width": width,
|
||||
"seed": seed,
|
||||
"rand_device": "cpu",
|
||||
"num_inference_steps": GLOBAL_STEPS,
|
||||
# The source graph applies ModelSamplingAuraFlow with shift=3. DiffSynth
|
||||
# expresses the same rational sigma shift as exp(mu), so mu=log(3).
|
||||
"exponential_shift_mu": math.log(3.0),
|
||||
"blockwise_controlnet_inputs": [controlnet_input],
|
||||
}
|
||||
|
||||
|
||||
def build_face_kwargs(crop: Image.Image, *, strength: float, seed: int | None) -> dict[str, Any]:
|
||||
"""Build the DiffSynth Z-Image face-detail call shape."""
|
||||
input_image = _resize_to_target(crop)
|
||||
width, height = input_image.size
|
||||
return {
|
||||
"prompt": _FACE_PROMPT,
|
||||
"negative_prompt": _FACE_NEGATIVE,
|
||||
"cfg_scale": FACE_CFG,
|
||||
"input_image": input_image,
|
||||
"denoising_strength": float(strength),
|
||||
"height": height,
|
||||
"width": width,
|
||||
"seed": seed,
|
||||
"rand_device": "cpu",
|
||||
"num_inference_steps": FACE_STEPS,
|
||||
}
|
||||
|
||||
|
||||
def composite_face(base: np.ndarray, detail: np.ndarray, mask: np.ndarray, *, feather: int = 10) -> np.ndarray:
|
||||
"""Feather ``detail`` into ``base`` while preserving every zero-mask pixel."""
|
||||
import cv2
|
||||
|
||||
if base.shape != detail.shape:
|
||||
raise ValueError("base and detail must have identical shapes")
|
||||
if mask.shape != base.shape[:2]:
|
||||
raise ValueError("mask must match the image height and width")
|
||||
|
||||
alpha = mask.astype(np.float32) / 255.0
|
||||
if feather > 0:
|
||||
sigma = max(0.1, float(feather) / 3.0)
|
||||
alpha = cv2.GaussianBlur(alpha, (0, 0), sigmaX=sigma, sigmaY=sigma)
|
||||
# Blurring may introduce tiny values far outside the intended mask. Keep an
|
||||
# explicit support dilation so pixels beyond the feather radius stay exact.
|
||||
support = cv2.dilate((mask > 0).astype(np.uint8), np.ones((2 * feather + 1, 2 * feather + 1), np.uint8))
|
||||
alpha *= support
|
||||
alpha = np.clip(alpha, 0.0, 1.0)[:, :, None]
|
||||
merged = base.astype(np.float32) * (1.0 - alpha) + detail.astype(np.float32) * alpha
|
||||
return np.clip(np.rint(merged), 0, 255).astype(np.uint8)
|
||||
|
||||
|
||||
def _expanded_box(
|
||||
box: tuple[int, int, int, int],
|
||||
image_size: tuple[int, int],
|
||||
*,
|
||||
factor: float = 2.5,
|
||||
) -> tuple[int, int, int, int]:
|
||||
"""Expand a face box around its center, matching the reference crop factor."""
|
||||
x1, y1, x2, y2 = box
|
||||
image_width, image_height = image_size
|
||||
center_x = (x1 + x2) / 2.0
|
||||
center_y = (y1 + y2) / 2.0
|
||||
width = max(1.0, (x2 - x1) * factor)
|
||||
height = max(1.0, (y2 - y1) * factor)
|
||||
return (
|
||||
max(0, round(center_x - width / 2.0)),
|
||||
max(0, round(center_y - height / 2.0)),
|
||||
min(image_width, round(center_x + width / 2.0)),
|
||||
min(image_height, round(center_y + height / 2.0)),
|
||||
)
|
||||
|
||||
|
||||
def _model_cache_dir() -> Path:
|
||||
root = os.environ.get("XDG_CACHE_HOME")
|
||||
base = Path(root) if root else Path.home() / ".cache"
|
||||
return base / "remove-ai-watermarks"
|
||||
|
||||
|
||||
def _yunet_model_path() -> Path:
|
||||
"""Download the small MIT-licensed YuNet ONNX model on first use."""
|
||||
model_path = _model_cache_dir() / YUNET_MODEL_NAME
|
||||
if model_path.exists() and hashlib.sha256(model_path.read_bytes()).hexdigest() == YUNET_MODEL_SHA256:
|
||||
return model_path
|
||||
|
||||
model_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
log.info("Downloading YuNet face detector: %s", YUNET_MODEL_URL)
|
||||
request = urllib.request.Request(
|
||||
YUNET_MODEL_URL,
|
||||
headers={"User-Agent": "remove-ai-watermarks"},
|
||||
)
|
||||
with urllib.request.urlopen(request, timeout=60) as response: # noqa: S310 - fixed HTTPS source
|
||||
payload = response.read()
|
||||
log.info(
|
||||
"YuNet download response: status=%s content_length=%s",
|
||||
getattr(response, "status", None),
|
||||
len(payload),
|
||||
)
|
||||
digest = hashlib.sha256(payload).hexdigest()
|
||||
if digest != YUNET_MODEL_SHA256:
|
||||
raise OSError(
|
||||
"YuNet download failed integrity verification: "
|
||||
f"expected {YUNET_MODEL_SHA256}, got {digest} ({len(payload)} bytes)"
|
||||
)
|
||||
with tempfile.NamedTemporaryFile(dir=model_path.parent, suffix=".onnx", delete=False) as handle:
|
||||
handle.write(payload)
|
||||
temporary = Path(handle.name)
|
||||
temporary.replace(model_path)
|
||||
return model_path
|
||||
|
||||
|
||||
def _nms_boxes(
|
||||
boxes: list[tuple[int, int, int, int]],
|
||||
scores: list[float],
|
||||
*,
|
||||
threshold: float = 0.3,
|
||||
) -> list[tuple[int, int, int, int]]:
|
||||
"""Apply OpenCV NMS to detections collected at multiple image scales."""
|
||||
if not boxes:
|
||||
return []
|
||||
import cv2
|
||||
|
||||
xywh = [[x1, y1, x2 - x1, y2 - y1] for x1, y1, x2, y2 in boxes]
|
||||
indices = cv2.dnn.NMSBoxes(xywh, scores, score_threshold=0.2, nms_threshold=threshold)
|
||||
if len(indices) == 0:
|
||||
return []
|
||||
return [boxes[int(index)] for index in np.asarray(indices).reshape(-1)]
|
||||
|
||||
|
||||
def detect_faces(image: Image.Image) -> list[tuple[int, int, int, int]]:
|
||||
"""Detect face boxes with YuNet at two scales for large and small faces."""
|
||||
import cv2
|
||||
|
||||
rgb = np.asarray(image.convert("RGB"))
|
||||
original_height, original_width = rgb.shape[:2]
|
||||
detections: list[tuple[int, int, int, int]] = []
|
||||
scores: list[float] = []
|
||||
model_path = _yunet_model_path()
|
||||
|
||||
for long_side in (640, 1280):
|
||||
scale = min(1.0, long_side / max(original_width, original_height))
|
||||
width = max(1, round(original_width * scale))
|
||||
height = max(1, round(original_height * scale))
|
||||
resized = cv2.resize(rgb, (width, height), interpolation=cv2.INTER_AREA) if scale < 1.0 else rgb
|
||||
bgr = cv2.cvtColor(resized, cv2.COLOR_RGB2BGR)
|
||||
detector = cv2.FaceDetectorYN.create(
|
||||
str(model_path),
|
||||
"",
|
||||
(width, height),
|
||||
YUNET_SCORE_THRESHOLD,
|
||||
0.3,
|
||||
5000,
|
||||
)
|
||||
_, rows = detector.detect(bgr)
|
||||
if rows is None:
|
||||
continue
|
||||
inverse = 1.0 / scale
|
||||
for row in rows:
|
||||
x, y, box_width, box_height = (float(value) for value in row[:4])
|
||||
x1 = max(0, round(x * inverse))
|
||||
y1 = max(0, round(y * inverse))
|
||||
x2 = min(original_width, round((x + box_width) * inverse))
|
||||
y2 = min(original_height, round((y + box_height) * inverse))
|
||||
if x2 <= x1 or y2 <= y1:
|
||||
continue
|
||||
detections.append((x1, y1, x2, y2))
|
||||
scores.append(float(row[-1]))
|
||||
if scale == 1.0:
|
||||
break
|
||||
return _nms_boxes(detections, scores)
|
||||
|
||||
|
||||
def _ellipse_masks(
|
||||
boxes: list[tuple[int, int, int, int]],
|
||||
image_size: tuple[int, int],
|
||||
) -> list[np.ndarray]:
|
||||
"""Safe fallback masks when SAM is unavailable."""
|
||||
import cv2
|
||||
|
||||
width, height = image_size
|
||||
masks: list[np.ndarray] = []
|
||||
for x1, y1, x2, y2 in boxes:
|
||||
mask = np.zeros((height, width), dtype=np.uint8)
|
||||
center = ((x1 + x2) // 2, (y1 + y2) // 2)
|
||||
axes = (max(1, int((x2 - x1) * 0.55)), max(1, int((y2 - y1) * 0.62)))
|
||||
cv2.ellipse(mask, center, axes, 0, 0, 360, 255, -1)
|
||||
masks.append(mask)
|
||||
return masks
|
||||
|
||||
|
||||
def _prepare_sam_inputs(inputs: Any, device: str, dtype: Any) -> Any:
|
||||
"""Move SAM inputs to the target device without casting geometric prompts."""
|
||||
prepared = inputs.to(device)
|
||||
if "pixel_values" in prepared:
|
||||
prepared["pixel_values"] = prepared["pixel_values"].to(dtype=dtype)
|
||||
return prepared
|
||||
|
||||
|
||||
def _sam_point_prompts(
|
||||
boxes: list[tuple[int, int, int, int]],
|
||||
) -> tuple[list[list[list[list[float]]]], list[list[list[int]]]]:
|
||||
"""Build Impact Pack's center-1 positive prompt for every face box."""
|
||||
points = [[[[(x1 + x2) / 2.0, (y1 + y2) / 2.0]] for x1, y1, x2, y2 in boxes]]
|
||||
labels = [[[1] for _box in boxes]]
|
||||
return points, labels
|
||||
|
||||
|
||||
def _clip_sam_masks_to_boxes(
|
||||
masks: list[np.ndarray],
|
||||
boxes: list[tuple[int, int, int, int]],
|
||||
image_size: tuple[int, int],
|
||||
) -> list[np.ndarray]:
|
||||
"""Match Impact Pack by intersecting each SAM mask with its detector box."""
|
||||
width, height = image_size
|
||||
clipped: list[np.ndarray] = []
|
||||
for mask, (x1, y1, x2, y2) in zip(masks, boxes, strict=True):
|
||||
box_mask = np.zeros((height, width), dtype=np.uint8)
|
||||
box_mask[max(0, y1) : min(height, y2), max(0, x1) : min(width, x2)] = 255
|
||||
clipped.append(np.bitwise_and(mask.astype(np.uint8), box_mask))
|
||||
return clipped
|
||||
|
||||
|
||||
def _select_sam_masks(
|
||||
masks: np.ndarray,
|
||||
scores: np.ndarray,
|
||||
*,
|
||||
threshold: float = 0.93,
|
||||
) -> list[np.ndarray]:
|
||||
"""Select and combine SAM proposals like Impact Pack's ``sub_threshold``."""
|
||||
mask_array = np.asarray(masks)
|
||||
score_array = np.asarray(scores)
|
||||
if mask_array.ndim == 5 and mask_array.shape[0] == 1:
|
||||
mask_array = mask_array[0]
|
||||
if score_array.ndim == 3 and score_array.shape[0] == 1:
|
||||
score_array = score_array[0]
|
||||
if mask_array.ndim == 3:
|
||||
mask_array = mask_array[:, None, :, :]
|
||||
if score_array.ndim == 1:
|
||||
score_array = score_array[:, None]
|
||||
if mask_array.ndim != 4 or score_array.ndim != 2:
|
||||
raise ValueError("SAM masks and scores have unexpected dimensions")
|
||||
if mask_array.shape[:2] != score_array.shape:
|
||||
raise ValueError("SAM mask proposals and IoU scores do not align")
|
||||
|
||||
selected_masks: list[np.ndarray] = []
|
||||
for candidates, candidate_scores in zip(mask_array, score_array, strict=True):
|
||||
selected = np.flatnonzero(candidate_scores >= threshold)
|
||||
if selected.size == 0:
|
||||
selected = np.asarray([int(np.argmax(candidate_scores))])
|
||||
combined = np.any(candidates[selected] > 0, axis=0)
|
||||
selected_masks.append(combined.astype(np.uint8) * 255)
|
||||
return selected_masks
|
||||
|
||||
|
||||
def _sam_outputs_to_numpy(masks: Any, scores: Any) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Convert SAM tensors through float32 because NumPy rejects bfloat16."""
|
||||
mask_array = masks.detach().float().cpu().numpy()
|
||||
score_array = scores.detach().float().cpu().numpy()
|
||||
return mask_array, score_array
|
||||
|
||||
|
||||
@dataclass
|
||||
class QwenZImagePipeline:
|
||||
"""Lazy runtime for the two-stage Qwen/Z-Image profile."""
|
||||
|
||||
device: str
|
||||
torch_dtype: Any
|
||||
hf_token: str | None = None
|
||||
progress_callback: Callable[[str], None] | None = None
|
||||
controlnet_conditioning_scale: float = GLOBAL_CONTROLNET_SCALE
|
||||
keep_face_models_on_device: bool | None = None
|
||||
cache_prompt_embeddings: bool = True
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._qwen_pipe: Any = None
|
||||
self._zimage_pipe: Any = None
|
||||
self._sam_model: Any = None
|
||||
self._sam_processor: Any = None
|
||||
|
||||
def _progress(self, message: str) -> None:
|
||||
if self.progress_callback is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
self.progress_callback(message)
|
||||
|
||||
def _require_cuda(self) -> None:
|
||||
if self.device != "cuda":
|
||||
raise RuntimeError(
|
||||
"The qwen-zimage pipeline is CUDA-only. Its Qwen-Image-2512 and "
|
||||
"Z-Image models do not fit the supported MPS path."
|
||||
)
|
||||
|
||||
def _vram_limit(self) -> float | None:
|
||||
import torch
|
||||
|
||||
with contextlib.suppress(Exception):
|
||||
return max(1.0, torch.cuda.mem_get_info("cuda")[1] / (1024**3) - 0.5)
|
||||
return None
|
||||
|
||||
def _keep_face_models_resident(self) -> bool:
|
||||
import torch
|
||||
|
||||
total_memory_gib = 0.0
|
||||
with contextlib.suppress(Exception):
|
||||
total_memory_gib = torch.cuda.get_device_properties("cuda").total_memory / (1024**3)
|
||||
return resolve_face_model_residency(
|
||||
self.keep_face_models_on_device,
|
||||
total_memory_gib=total_memory_gib,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _qwen_vram_config() -> dict[str, Any]:
|
||||
import torch
|
||||
|
||||
return {
|
||||
"offload_dtype": "disk",
|
||||
"offload_device": "disk",
|
||||
"onload_dtype": torch.float8_e4m3fn,
|
||||
"onload_device": "cpu",
|
||||
"preparing_dtype": torch.float8_e4m3fn,
|
||||
"preparing_device": "cuda",
|
||||
"computation_dtype": torch.bfloat16,
|
||||
"computation_device": "cuda",
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _zimage_vram_config() -> dict[str, Any]:
|
||||
import torch
|
||||
|
||||
return {
|
||||
"offload_dtype": torch.bfloat16,
|
||||
"offload_device": "cpu",
|
||||
"onload_dtype": torch.bfloat16,
|
||||
"onload_device": "cpu",
|
||||
"preparing_dtype": torch.bfloat16,
|
||||
"preparing_device": "cuda",
|
||||
"computation_dtype": torch.bfloat16,
|
||||
"computation_device": "cuda",
|
||||
}
|
||||
|
||||
def _load_qwen(self) -> Any:
|
||||
if self._qwen_pipe is not None:
|
||||
return self._qwen_pipe
|
||||
self._require_cuda()
|
||||
os.environ.setdefault("DIFFSYNTH_DOWNLOAD_SOURCE", "huggingface")
|
||||
if self.hf_token:
|
||||
os.environ.setdefault("HF_TOKEN", self.hf_token)
|
||||
try:
|
||||
from diffsynth.pipelines.qwen_image import ControlNetInput, ModelConfig, QwenImagePipeline
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"The qwen-zimage pipeline needs the optional dependency group. "
|
||||
"Install: pip install 'remove-ai-watermarks[qwen-zimage]'"
|
||||
) from exc
|
||||
|
||||
self._progress("Loading Qwen-Image-2512, Lightning LoRA, and Canny ControlNet...")
|
||||
config = self._qwen_vram_config()
|
||||
model_configs = [
|
||||
ModelConfig(
|
||||
model_id=QWEN_IMAGE_2512_MODEL_ID,
|
||||
origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors",
|
||||
**config,
|
||||
),
|
||||
ModelConfig(
|
||||
model_id=QWEN_IMAGE_2512_MODEL_ID,
|
||||
origin_file_pattern="text_encoder/model*.safetensors",
|
||||
**config,
|
||||
),
|
||||
ModelConfig(
|
||||
model_id=QWEN_IMAGE_2512_MODEL_ID,
|
||||
origin_file_pattern="vae/diffusion_pytorch_model.safetensors",
|
||||
**config,
|
||||
),
|
||||
ModelConfig(
|
||||
model_id=QWEN_CANNY_CONTROLNET_MODEL_ID,
|
||||
origin_file_pattern="model.safetensors",
|
||||
**config,
|
||||
),
|
||||
]
|
||||
pipe = QwenImagePipeline.from_pretrained(
|
||||
torch_dtype=self.torch_dtype,
|
||||
device=self.device,
|
||||
model_configs=model_configs,
|
||||
tokenizer_config=ModelConfig(
|
||||
model_id=QWEN_IMAGE_2512_MODEL_ID,
|
||||
origin_file_pattern="tokenizer/",
|
||||
),
|
||||
vram_limit=self._vram_limit(),
|
||||
)
|
||||
from diffsynth.diffusion import FlowMatchScheduler
|
||||
|
||||
# Avoid the base Qwen scheduler's terminal rescale for the distilled LoRA.
|
||||
# With exponential_shift_mu=log(3), this is the closest DiffSynth equivalent
|
||||
# of the source graph's four-step sgm_uniform + AuraFlow shift 3.
|
||||
pipe.scheduler = FlowMatchScheduler("Qwen-Image-Lightning")
|
||||
lightning = ModelConfig(
|
||||
model_id=QWEN_LIGHTNING_MODEL_ID,
|
||||
origin_file_pattern=QWEN_LIGHTNING_PATTERN,
|
||||
)
|
||||
pipe.load_lora(pipe.dit, lightning, alpha=0.8)
|
||||
if self.cache_prompt_embeddings:
|
||||
_cache_static_prompt_embeddings(
|
||||
pipe,
|
||||
("prompt_emb", "prompt_emb_mask"),
|
||||
)
|
||||
self._qwen_pipe = (pipe, ControlNetInput)
|
||||
return self._qwen_pipe
|
||||
|
||||
def _load_zimage(self) -> Any:
|
||||
if self._zimage_pipe is not None:
|
||||
return self._zimage_pipe
|
||||
self._require_cuda()
|
||||
os.environ.setdefault("DIFFSYNTH_DOWNLOAD_SOURCE", "huggingface")
|
||||
if self.hf_token:
|
||||
os.environ.setdefault("HF_TOKEN", self.hf_token)
|
||||
try:
|
||||
from diffsynth.pipelines.z_image import ModelConfig, ZImagePipeline
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"The qwen-zimage pipeline needs the optional dependency group. "
|
||||
"Install: pip install 'remove-ai-watermarks[qwen-zimage]'"
|
||||
) from exc
|
||||
|
||||
self._progress("Loading Z-Image Turbo face-detail model...")
|
||||
keep_on_device = self._keep_face_models_resident()
|
||||
config = self._zimage_vram_config()
|
||||
pipe = ZImagePipeline.from_pretrained(
|
||||
torch_dtype=self.torch_dtype,
|
||||
device=self.device,
|
||||
model_configs=[
|
||||
ModelConfig(
|
||||
model_id=ZIMAGE_TURBO_MODEL_ID,
|
||||
origin_file_pattern="transformer/*.safetensors",
|
||||
**config,
|
||||
),
|
||||
ModelConfig(
|
||||
model_id=ZIMAGE_TURBO_MODEL_ID,
|
||||
origin_file_pattern="text_encoder/*.safetensors",
|
||||
**config,
|
||||
),
|
||||
ModelConfig(
|
||||
model_id=ZIMAGE_TURBO_MODEL_ID,
|
||||
origin_file_pattern="vae/diffusion_pytorch_model.safetensors",
|
||||
**config,
|
||||
),
|
||||
],
|
||||
tokenizer_config=ModelConfig(
|
||||
model_id=ZIMAGE_TURBO_MODEL_ID,
|
||||
origin_file_pattern="tokenizer/",
|
||||
),
|
||||
vram_limit=self._vram_limit(),
|
||||
)
|
||||
if keep_on_device:
|
||||
_pin_vram_managed_models(pipe)
|
||||
if self.cache_prompt_embeddings:
|
||||
_cache_static_prompt_embeddings(pipe, ("prompt_embeds",))
|
||||
self._zimage_pipe = pipe
|
||||
return pipe
|
||||
|
||||
def _load_sam(self) -> tuple[Any, Any]:
|
||||
if self._sam_model is not None and self._sam_processor is not None:
|
||||
return self._sam_model, self._sam_processor
|
||||
self._progress("Loading SAM face-mask model...")
|
||||
try:
|
||||
from transformers import AutoModelForMaskGeneration, AutoProcessor
|
||||
except ImportError as exc:
|
||||
raise ImportError("SAM needs transformers and torchvision from the qwen-zimage extra.") from exc
|
||||
kwargs: dict[str, Any] = {}
|
||||
if self.hf_token:
|
||||
kwargs["token"] = self.hf_token
|
||||
processor = AutoProcessor.from_pretrained(SAM_MODEL_ID, **kwargs)
|
||||
model = AutoModelForMaskGeneration.from_pretrained(
|
||||
SAM_MODEL_ID,
|
||||
torch_dtype=self.torch_dtype,
|
||||
**kwargs,
|
||||
).to(self.device)
|
||||
model.eval()
|
||||
self._sam_model = model
|
||||
self._sam_processor = processor
|
||||
return model, processor
|
||||
|
||||
def _sam_masks(
|
||||
self,
|
||||
image: Image.Image,
|
||||
boxes: list[tuple[int, int, int, int]],
|
||||
) -> list[np.ndarray]:
|
||||
if not boxes:
|
||||
return []
|
||||
import torch
|
||||
|
||||
try:
|
||||
model, processor = self._load_sam()
|
||||
input_points, input_labels = _sam_point_prompts(boxes)
|
||||
inputs = processor(
|
||||
images=image,
|
||||
input_boxes=[[list(box) for box in boxes]],
|
||||
input_points=input_points,
|
||||
input_labels=input_labels,
|
||||
return_tensors="pt",
|
||||
)
|
||||
original_sizes = inputs["original_sizes"].clone()
|
||||
reshaped_sizes = inputs["reshaped_input_sizes"].clone()
|
||||
inputs = _prepare_sam_inputs(inputs, self.device, self.torch_dtype)
|
||||
with torch.inference_mode():
|
||||
outputs = model(**inputs, multimask_output=True)
|
||||
processed = processor.post_process_masks(
|
||||
outputs.pred_masks.detach().cpu(),
|
||||
original_sizes,
|
||||
reshaped_sizes,
|
||||
)[0]
|
||||
mask_array, score_array = _sam_outputs_to_numpy(processed, outputs.iou_scores)
|
||||
binary_masks = _select_sam_masks(
|
||||
mask_array,
|
||||
score_array,
|
||||
)
|
||||
return _clip_sam_masks_to_boxes(binary_masks, boxes, image.size)
|
||||
except Exception as exc:
|
||||
log.warning("SAM face-mask refinement failed (%s); using box-derived ellipse masks", exc)
|
||||
return _ellipse_masks(boxes, image.size)
|
||||
|
||||
def preload(self) -> None:
|
||||
"""Eagerly load both diffusion stages and the face segmentation model."""
|
||||
self._load_qwen()
|
||||
self._load_zimage()
|
||||
self._load_sam()
|
||||
_yunet_model_path()
|
||||
|
||||
def _run_global(self, image: Image.Image, strength: float, seed: int | None) -> Image.Image:
|
||||
pipe, controlnet_input_cls = self._load_qwen()
|
||||
input_image = _resize_to_target(image)
|
||||
control = build_canny_control_image(input_image)
|
||||
control_input = controlnet_input_cls(
|
||||
image=control,
|
||||
scale=float(self.controlnet_conditioning_scale),
|
||||
)
|
||||
self._progress(f"Running Qwen-Image-2512 Canny pass: strength={strength:.4f}, steps={GLOBAL_STEPS}...")
|
||||
result = pipe(
|
||||
**build_global_kwargs(
|
||||
input_image,
|
||||
strength=strength,
|
||||
seed=seed,
|
||||
controlnet_input=control_input,
|
||||
)
|
||||
)
|
||||
if result.size != image.size:
|
||||
result = result.resize(image.size, Image.Resampling.LANCZOS)
|
||||
return result.convert("RGB")
|
||||
|
||||
@staticmethod
|
||||
def _detail_size(
|
||||
crop_size: tuple[int, int],
|
||||
face_size: tuple[int, int],
|
||||
) -> tuple[int, int]:
|
||||
"""Scale a crop toward a 768px face guide while capping it at 1024px."""
|
||||
crop_width, crop_height = crop_size
|
||||
face_width, face_height = face_size
|
||||
scale_for_face = 768.0 / max(1, max(face_width, face_height))
|
||||
scale_for_crop = 1024.0 / max(1, max(crop_width, crop_height))
|
||||
scale = min(scale_for_face, scale_for_crop)
|
||||
# Never shrink below the crop's current size unless the 1024 cap requires it.
|
||||
if max(crop_width, crop_height) <= 1024:
|
||||
scale = max(1.0, scale)
|
||||
width = max(16, round(crop_width * scale / 16.0) * 16)
|
||||
height = max(16, round(crop_height * scale / 16.0) * 16)
|
||||
return width, height
|
||||
|
||||
def _run_faces(
|
||||
self,
|
||||
original: Image.Image,
|
||||
global_result: Image.Image,
|
||||
boxes: list[tuple[int, int, int, int]],
|
||||
masks: list[np.ndarray],
|
||||
*,
|
||||
strength: float,
|
||||
seed: int | None,
|
||||
) -> Image.Image:
|
||||
if not boxes:
|
||||
return global_result
|
||||
pipe = self._load_zimage()
|
||||
base = np.asarray(global_result.convert("RGB")).copy()
|
||||
source = np.asarray(original.convert("RGB"))
|
||||
detail_seed = None if seed is None else seed + 1
|
||||
|
||||
for index, (box, mask) in enumerate(zip(boxes, masks, strict=True), start=1):
|
||||
crop_box = _expanded_box(box, original.size)
|
||||
cx1, cy1, cx2, cy2 = crop_box
|
||||
crop_source = source[cy1:cy2, cx1:cx2]
|
||||
crop_mask = mask[cy1:cy2, cx1:cx2]
|
||||
if crop_source.size == 0 or not np.any(crop_mask):
|
||||
continue
|
||||
|
||||
face_width = box[2] - box[0]
|
||||
face_height = box[3] - box[1]
|
||||
process_size = self._detail_size((cx2 - cx1, cy2 - cy1), (face_width, face_height))
|
||||
crop_image = Image.fromarray(crop_source).resize(process_size, Image.Resampling.LANCZOS)
|
||||
self._progress(
|
||||
f"Regenerating face {index}/{len(boxes)} with Z-Image: strength={strength:.4f}, steps={FACE_STEPS}..."
|
||||
)
|
||||
detailed = pipe(**build_face_kwargs(crop_image, strength=strength, seed=detail_seed))
|
||||
detailed = detailed.convert("RGB").resize((cx2 - cx1, cy2 - cy1), Image.Resampling.LANCZOS)
|
||||
|
||||
base_crop = base[cy1:cy2, cx1:cx2]
|
||||
base[cy1:cy2, cx1:cx2] = composite_face(
|
||||
base_crop,
|
||||
np.asarray(detailed),
|
||||
crop_mask,
|
||||
feather=10,
|
||||
)
|
||||
return Image.fromarray(base)
|
||||
|
||||
def run(
|
||||
self,
|
||||
image: Image.Image,
|
||||
*,
|
||||
strength: float | None,
|
||||
seed: int | None,
|
||||
tile: bool = False,
|
||||
tile_size: int = 1024,
|
||||
tile_overlap: int = 128,
|
||||
) -> Image.Image:
|
||||
"""Execute global regeneration and masked face repair."""
|
||||
self._require_cuda()
|
||||
seed = resolve_seed(seed, "qwen-zimage")
|
||||
global_strength = (
|
||||
resolution_adaptive_denoise(image.width, image.height) if strength is None else float(strength)
|
||||
)
|
||||
if tile and max(image.size) > tile_size:
|
||||
from remove_ai_watermarks.noai.tiling import run_tiled
|
||||
|
||||
global_result = run_tiled(
|
||||
lambda tile_image: self._run_global(tile_image, global_strength, seed),
|
||||
image,
|
||||
tile_size,
|
||||
tile_overlap,
|
||||
self._progress,
|
||||
)
|
||||
else:
|
||||
global_result = self._run_global(image, global_strength, seed)
|
||||
|
||||
self._progress("Detecting faces on the original image...")
|
||||
boxes = detect_faces(image)
|
||||
if not boxes:
|
||||
self._progress("No faces detected; keeping the Qwen global result.")
|
||||
return global_result
|
||||
masks = self._sam_masks(image, boxes)
|
||||
face_strength = largest_face_denoise(boxes, image.size)
|
||||
return self._run_faces(
|
||||
image,
|
||||
global_result,
|
||||
boxes,
|
||||
masks,
|
||||
strength=face_strength,
|
||||
seed=seed,
|
||||
)
|
||||
@@ -31,7 +31,11 @@ QWEN_MODEL_ID = "Qwen/Qwen-Image"
|
||||
# profile is ``sdxl``; ``default`` is kept as an accepted alias (it was the profile's
|
||||
# name before ``controlnet`` became the default-selected pipeline, 2026-06-09).
|
||||
SDXL_PROFILE = "sdxl"
|
||||
_PROFILE_ALIASES = {"default": SDXL_PROFILE}
|
||||
QWEN_ZIMAGE_PROFILE = "qwen-zimage"
|
||||
_PROFILE_ALIASES = {
|
||||
"default": SDXL_PROFILE,
|
||||
"qwen_zimage": QWEN_ZIMAGE_PROFILE,
|
||||
}
|
||||
|
||||
|
||||
def normalize_profile(profile: str) -> str:
|
||||
@@ -40,6 +44,24 @@ def normalize_profile(profile: str) -> str:
|
||||
return _PROFILE_ALIASES.get(normalized, normalized)
|
||||
|
||||
|
||||
def resolve_steps(num_inference_steps: int | None, pipeline: str) -> int:
|
||||
"""Resolve a profile-specific step default while preserving explicit values.
|
||||
|
||||
The Lightning LoRA in ``qwen-zimage`` is distilled for four steps. Existing
|
||||
SDXL and Qwen profiles keep the long-standing 50-step CLI default.
|
||||
"""
|
||||
if num_inference_steps is not None:
|
||||
return num_inference_steps
|
||||
return 4 if normalize_profile(pipeline) == QWEN_ZIMAGE_PROFILE else 50
|
||||
|
||||
|
||||
def resolve_seed(seed: int | None, pipeline: str) -> int | None:
|
||||
"""Keep the oracle-verified qwen-zimage profile deterministic by default."""
|
||||
if seed is not None:
|
||||
return seed
|
||||
return 0 if normalize_profile(pipeline) == QWEN_ZIMAGE_PROFILE else None
|
||||
|
||||
|
||||
# The SDXL-native canny ControlNet used by the ``controlnet`` pipeline. The
|
||||
# ControlNet is an add-on to the SDXL base checkpoint (DEFAULT_MODEL_ID), not a
|
||||
# separate base model, so both the ``sdxl`` and ``controlnet`` profiles load the
|
||||
@@ -113,7 +135,8 @@ def strength_default_help() -> str:
|
||||
"""
|
||||
return (
|
||||
f"vendor-adaptive (OpenAI {OPENAI_STRENGTH} / Google {GEMINI_STRENGTH} / "
|
||||
f"unknown {UNKNOWN_STRENGTH}, from the C2PA issuer; same ladder for both pipelines)"
|
||||
f"unknown {UNKNOWN_STRENGTH}, from the C2PA issuer; qwen-zimage instead uses "
|
||||
"resolution-adaptive denoise)"
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
"""Watermark removal using diffusion model regeneration attack.
|
||||
|
||||
Three pipelines (selected by the explicit ``pipeline`` ctor arg):
|
||||
Four pipelines (selected by the explicit ``pipeline`` ctor arg):
|
||||
|
||||
0. ``qwen`` -- Qwen-Image (20B MMDiT, Apache-2.0) img2img. The scrub still comes from
|
||||
0. ``qwen-zimage`` -- Qwen-Image-2512 Lightning + DiffSynth Canny regenerates the
|
||||
frame, then SAM-masked Z-Image Turbo regenerates original face crops and feathers
|
||||
them into the global result. CUDA-only and installed through its own optional extra.
|
||||
1. ``qwen`` -- Qwen-Image (20B MMDiT, Apache-2.0) img2img. The scrub still comes from
|
||||
the img2img ``strength``; Qwen preserves text (incl. CJK) and structure markedly
|
||||
better than SDXL at the scrub floor, so it over-regenerates real photos far less.
|
||||
CUDA/cloud-class (does not fit MPS). See ``watermark_profiles`` for the certified
|
||||
@@ -11,14 +14,14 @@ Three pipelines (selected by the explicit ``pipeline`` ctor arg):
|
||||
is the release gate, not sweeping seeds.)
|
||||
|
||||
Two SDXL pipelines:
|
||||
1. ``controlnet`` (DEFAULT) -- SDXL img2img with a canny ControlNet. The watermark
|
||||
2. ``controlnet`` (DEFAULT) -- SDXL img2img with a canny ControlNet. The watermark
|
||||
REMOVAL still comes from the img2img regeneration (``strength``); the ControlNet
|
||||
only PRESERVES structure (text/faces) by conditioning on the edge map. No original
|
||||
pixels are ever copied or frozen. Because the edge map keeps the regeneration
|
||||
closer to the original, it needs a higher ``strength`` floor than ``default`` to
|
||||
destroy SynthID (the certified controlnet ladder; see ``watermark_profiles``).
|
||||
``controlnet_conditioning_scale`` is the preservation knob.
|
||||
2. ``default`` -- plain SDXL img2img. Partial-noise regeneration scrubs the
|
||||
3. ``default`` -- plain SDXL img2img. Partial-noise regeneration scrubs the
|
||||
invisible watermark; ``strength`` controls how much is regenerated. Lighter (no
|
||||
ControlNet weights), but at the low default strength it leaves SynthID on
|
||||
flat-graphic content -- use it for inputs without text/faces.
|
||||
@@ -47,7 +50,10 @@ from remove_ai_watermarks.noai.watermark_profiles import (
|
||||
DEFAULT_MODEL_ID,
|
||||
DEFAULT_STRENGTH,
|
||||
QWEN_MODEL_ID,
|
||||
QWEN_ZIMAGE_PROFILE,
|
||||
normalize_profile,
|
||||
resolve_seed,
|
||||
resolve_steps,
|
||||
resolve_strength,
|
||||
viable_steps,
|
||||
)
|
||||
@@ -388,6 +394,13 @@ class WatermarkRemover:
|
||||
# the legacy "default" alias resolves to "sdxl".
|
||||
self.model_profile = normalize_profile(pipeline)
|
||||
self.controlnet_conditioning_scale = controlnet_conditioning_scale
|
||||
if self.model_profile == QWEN_ZIMAGE_PROFILE and self.model_id != self.DEFAULT_MODEL_ID:
|
||||
raise ValueError(
|
||||
"The qwen-zimage pipeline uses a fixed Qwen-Image-2512 + Z-Image model stack; "
|
||||
"--model is not supported for this profile."
|
||||
)
|
||||
if self.model_profile == QWEN_ZIMAGE_PROFILE:
|
||||
self.model_id = "Qwen/Qwen-Image-2512 + Tongyi-MAI/Z-Image-Turbo"
|
||||
|
||||
if not is_watermark_removal_available():
|
||||
_ensure_watermark_deps()
|
||||
@@ -399,8 +412,8 @@ class WatermarkRemover:
|
||||
if torch_dtype is None:
|
||||
if self.device == "cpu" or self.device == "mps":
|
||||
self.torch_dtype = torch.float32 # type: ignore
|
||||
elif self.model_profile == "qwen":
|
||||
# Qwen-Image is published in bf16; fp16 risks overflow on the 20B MMDiT.
|
||||
elif self.model_profile in {"qwen", QWEN_ZIMAGE_PROFILE}:
|
||||
# Qwen-Image and Z-Image are published in bf16; fp16 risks overflow.
|
||||
# cuda/xpu-only by construction: the cpu/mps guard above already forced
|
||||
# fp32, and the 20B model does not fit MPS anyway.
|
||||
self.torch_dtype = torch.bfloat16 # type: ignore
|
||||
@@ -412,6 +425,7 @@ class WatermarkRemover:
|
||||
self._pipeline: AutoImg2ImgPipeline | None = None
|
||||
self._controlnet_pipeline: Any = None
|
||||
self._qwen_pipeline: Any = None
|
||||
self._qwen_zimage_pipeline: Any = None
|
||||
self._progress_callback = progress_callback
|
||||
self.hf_token: str | None = hf_token or os.environ.get("HF_TOKEN")
|
||||
|
||||
@@ -426,7 +440,9 @@ class WatermarkRemover:
|
||||
|
||||
def preload(self) -> None:
|
||||
"""Eagerly load the pipeline so download progress bars are visible."""
|
||||
if self.model_profile == "qwen":
|
||||
if self.model_profile == QWEN_ZIMAGE_PROFILE:
|
||||
self._load_qwen_zimage_pipeline().preload()
|
||||
elif self.model_profile == "qwen":
|
||||
self._load_qwen_pipeline()
|
||||
elif self.model_profile == "controlnet":
|
||||
self._load_controlnet_pipeline()
|
||||
@@ -611,6 +627,20 @@ class WatermarkRemover:
|
||||
|
||||
return self._qwen_pipeline
|
||||
|
||||
def _load_qwen_zimage_pipeline(self) -> Any:
|
||||
"""Load the two-stage Qwen-Image-2512 + Z-Image runtime lazily."""
|
||||
if self._qwen_zimage_pipeline is None:
|
||||
from remove_ai_watermarks.noai.qwen_zimage_pipeline import QwenZImagePipeline
|
||||
|
||||
self._qwen_zimage_pipeline = QwenZImagePipeline(
|
||||
device=self.device,
|
||||
torch_dtype=self.torch_dtype,
|
||||
hf_token=self.hf_token,
|
||||
progress_callback=self._progress_callback,
|
||||
controlnet_conditioning_scale=self.controlnet_conditioning_scale,
|
||||
)
|
||||
return self._qwen_zimage_pipeline
|
||||
|
||||
# ── Core removal ─────────────────────────────────────────────────
|
||||
|
||||
def remove_watermark(
|
||||
@@ -618,7 +648,7 @@ class WatermarkRemover:
|
||||
image_path: Path,
|
||||
output_path: Path | None = None,
|
||||
strength: float | None = None,
|
||||
num_inference_steps: int = 50,
|
||||
num_inference_steps: int | None = None,
|
||||
guidance_scale: float | None = None,
|
||||
seed: int | None = None,
|
||||
vendor: str | None = None,
|
||||
@@ -637,7 +667,8 @@ class WatermarkRemover:
|
||||
default (see ``vendor``).
|
||||
num_inference_steps: Number of denoising steps.
|
||||
guidance_scale: Classifier-free guidance scale.
|
||||
seed: Random seed for reproducibility.
|
||||
seed: Random seed for reproducibility. None resolves to 0 for
|
||||
qwen-zimage and stays random for the other profiles.
|
||||
vendor: SynthID vendor (``"openai"`` / ``"google"`` / None) used to pick the
|
||||
default strength when ``strength`` is None. Detect it from the ORIGINAL
|
||||
input with ``watermark_profiles.vendor_for_strength`` before processing
|
||||
@@ -647,7 +678,7 @@ class WatermarkRemover:
|
||||
The lossless alternative to a ``--max-resolution`` downscale for large
|
||||
inputs that OOM on MPS/GPU (issue #10). Only engages when the long side
|
||||
exceeds ``tile_size``; smaller images run a single pass unchanged.
|
||||
tile_size: Tile dimension in px (default 1024, SDXL's training size).
|
||||
tile_size: Tile dimension in px (default 1024).
|
||||
tile_overlap: Overlap between adjacent tiles in px (default 128), feather-
|
||||
blended so there is no visible seam.
|
||||
region: Restrict the regeneration to the AI-composited box ``(x, y, w, h)``
|
||||
@@ -673,19 +704,27 @@ class WatermarkRemover:
|
||||
if output_path is None:
|
||||
output_path = image_path
|
||||
|
||||
strength = resolve_strength(strength, vendor, self.model_profile)
|
||||
|
||||
if not 0.0 <= strength <= 1.0:
|
||||
raise ValueError(f"Strength must be between 0.0 and 1.0, got {strength}")
|
||||
|
||||
if guidance_scale is None:
|
||||
guidance_scale = 7.5
|
||||
|
||||
self._set_progress("Loading and preprocessing input image...")
|
||||
init_image = Image.open(image_path).convert("RGB")
|
||||
w, h = init_image.size
|
||||
self._set_progress(f"Image loaded: {w}x{h}px | Model: {self.model_id}")
|
||||
|
||||
if self.model_profile == QWEN_ZIMAGE_PROFILE:
|
||||
from remove_ai_watermarks.noai.qwen_zimage_pipeline import resolution_adaptive_denoise
|
||||
|
||||
strength = strength if strength is not None else resolution_adaptive_denoise(w, h)
|
||||
else:
|
||||
strength = resolve_strength(strength, vendor, self.model_profile)
|
||||
seed = resolve_seed(seed, self.model_profile)
|
||||
if not 0.0 <= strength <= 1.0:
|
||||
raise ValueError(f"Strength must be between 0.0 and 1.0, got {strength}")
|
||||
|
||||
num_inference_steps = resolve_steps(num_inference_steps, self.model_profile)
|
||||
if guidance_scale is None:
|
||||
guidance_scale = 1.0 if self.model_profile == QWEN_ZIMAGE_PROFILE else 7.5
|
||||
elif self.model_profile == QWEN_ZIMAGE_PROFILE and guidance_scale != 1.0:
|
||||
raise ValueError("The qwen-zimage profile fixes both diffusion stages at CFG 1.0.")
|
||||
|
||||
generator = None
|
||||
if seed is not None and _HAS_TORCH:
|
||||
self._set_progress(f"Setting reproducible seed: {seed}")
|
||||
@@ -695,17 +734,26 @@ class WatermarkRemover:
|
||||
# inside attention with an opaque reshape error, so raise it to the minimum that
|
||||
# denoises. Must be applied to the value HANDED TO THE PIPELINE -- the old
|
||||
# max(1, ...) below only clamped the number in the log line.
|
||||
adjusted = viable_steps(num_inference_steps, strength)
|
||||
if adjusted != num_inference_steps:
|
||||
logger.warning(
|
||||
"steps=%s at strength=%s denoises 0 steps and would crash; using steps=%s (1 effective)",
|
||||
num_inference_steps,
|
||||
strength,
|
||||
adjusted,
|
||||
)
|
||||
num_inference_steps = adjusted
|
||||
if self.model_profile != QWEN_ZIMAGE_PROFILE:
|
||||
adjusted = viable_steps(num_inference_steps, strength)
|
||||
if adjusted != num_inference_steps:
|
||||
logger.warning(
|
||||
"steps=%s at strength=%s denoises 0 steps and would crash; using steps=%s (1 effective)",
|
||||
num_inference_steps,
|
||||
strength,
|
||||
adjusted,
|
||||
)
|
||||
num_inference_steps = adjusted
|
||||
elif num_inference_steps != 4:
|
||||
raise ValueError("The qwen-zimage profile uses the 4-step Lightning LoRA, so --steps must be 4.")
|
||||
|
||||
effective_steps = max(1, int(num_inference_steps * strength))
|
||||
# DiffSynth keeps all timesteps and compresses their sigma range for a low
|
||||
# denoise value. Diffusers instead truncates the schedule by strength.
|
||||
effective_steps = (
|
||||
num_inference_steps
|
||||
if self.model_profile == QWEN_ZIMAGE_PROFILE
|
||||
else max(1, int(num_inference_steps * strength))
|
||||
)
|
||||
self._set_progress(
|
||||
f"Config: strength={strength}, steps={num_inference_steps} "
|
||||
f"(~{effective_steps} effective), guidance={guidance_scale}, device={self.device}"
|
||||
@@ -721,6 +769,17 @@ class WatermarkRemover:
|
||||
return self._run_img2img(img, strength, num_inference_steps, guidance_scale, generator)
|
||||
|
||||
def _generate() -> Image.Image:
|
||||
# qwen-zimage owns its global-only tiling because its face stage must run
|
||||
# once after the tiles are blended. Other profiles tile their whole pass.
|
||||
if self.model_profile == QWEN_ZIMAGE_PROFILE:
|
||||
return self._run_qwen_zimage(
|
||||
init_image,
|
||||
strength,
|
||||
seed,
|
||||
tile=tile,
|
||||
tile_size=tile_size,
|
||||
tile_overlap=tile_overlap,
|
||||
)
|
||||
# Tile only when asked AND the image is larger than one tile; otherwise a
|
||||
# single full-image pass (tiling a sub-tile image is pure overhead).
|
||||
if tile and max(init_image.size) > tile_size:
|
||||
@@ -742,6 +801,7 @@ class WatermarkRemover:
|
||||
self.torch_dtype = torch.float32
|
||||
self._pipeline = None
|
||||
self._controlnet_pipeline = None
|
||||
self._qwen_zimage_pipeline = None
|
||||
cleaned_image = _generate()
|
||||
|
||||
# Region-targeted regeneration for AI-enhanced composites: keep the real photo
|
||||
@@ -941,6 +1001,27 @@ class WatermarkRemover:
|
||||
result = pipeline(**kwargs)
|
||||
return result.images[0]
|
||||
|
||||
def _run_qwen_zimage(
|
||||
self,
|
||||
init_image: Image.Image,
|
||||
strength: float,
|
||||
seed: int | None,
|
||||
*,
|
||||
tile: bool = False,
|
||||
tile_size: int = 1024,
|
||||
tile_overlap: int = 128,
|
||||
) -> Image.Image:
|
||||
"""Run the Qwen 2512 Canny pass and masked Z-Image face repair."""
|
||||
pipeline = self._load_qwen_zimage_pipeline()
|
||||
return pipeline.run(
|
||||
init_image,
|
||||
strength=strength,
|
||||
seed=seed,
|
||||
tile=tile,
|
||||
tile_size=tile_size,
|
||||
tile_overlap=tile_overlap,
|
||||
)
|
||||
|
||||
# ── Batch ────────────────────────────────────────────────────────
|
||||
|
||||
def remove_watermark_batch(
|
||||
@@ -948,7 +1029,7 @@ class WatermarkRemover:
|
||||
input_dir: Path,
|
||||
output_dir: Path,
|
||||
strength: float | None = None,
|
||||
num_inference_steps: int = 50,
|
||||
num_inference_steps: int | None = None,
|
||||
extensions: tuple[str, ...] = (".png", ".jpg", ".jpeg", ".webp"),
|
||||
) -> list[Path]:
|
||||
"""Remove watermarks from all images in a directory."""
|
||||
|
||||
Reference in New Issue
Block a user