From e42b7e9d6a10d8e32714abecaa8e691df7b5a3cb Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Sun, 31 May 2026 15:21:29 -0700 Subject: [PATCH] refactor(cli): plain-text console output; drop rich; quiet transformers cli.py now emits plain ASCII through a small click.echo shim (_Console / _Table / _Progress) instead of rich: no colors, markup tags, panels, progress bar, or Unicode glyphs (Warning: / -> / ... and dropped checkmark/cross marks). identify and metadata tables render as indented plain lines. - drop rich from dependencies (pyproject.toml + uv.lock) - __init__: set TRANSFORMERS_VERBOSITY=error (setdefault) plus a warnings filter so the transformers Siglip2ImageProcessorFast deprecation no longer prints at CLI startup (it fires from the eager noai import) - TestGpuHintMarkup: the [gpu] hint is now printed verbatim; docstring updated - CLAUDE.md: replace the obsolete rich-markup lesson, note the verbosity fix Co-Authored-By: Claude Opus 4.8 --- CLAUDE.md | 2 +- pyproject.toml | 1 - src/remove_ai_watermarks/__init__.py | 11 ++ src/remove_ai_watermarks/cli.py | 269 +++++++++++++++++---------- tests/test_cli.py | 4 +- uv.lock | 2 - 6 files changed, 185 insertions(+), 104 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index da63e83..b7e1a64 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -76,7 +76,7 @@ Who embeds what, and whether it is locally detectable (so we know which gaps are ## Known limitations -- `invisible` pipeline processes at **native resolution by default** (`max_resolution=0`), matching the hosted raiw.cc backend (fal fast-sdxl, no pre-downscale). The old forced downscale-to-1024 -> upscale-back round-trip was the main quality loss (issue #10) and is gone; at strength ~0.05 SDXL img2img does not need the ~1024 downscale. `--max-resolution N` re-introduces an opt-in long-side cap purely to bound GPU/MPS memory on very large inputs (it reintroduces the lossy round-trip). For huge images that OOM at native, tile-based diffusion is still the proper long-term fix. **Concrete MPS data point (verified 2026-05-25 on a 1254x1254 gpt-image SDXL run, fp32, 20 GB MPS ceiling):** native res OOMs at the *UNet* step (peak ~17 GiB), not only the VAE decode, and the auto-fallback in `img2img_runner` reloads on CPU and finishes (slow, ~13 min) -- the output is still weight-identical and defeats SynthID, so "looks hung/crashed" on Mac is usually this CPU fallback, not a pipeline error. Adding `enable_vae_tiling()` alone does NOT prevent it (the peak is the UNet, not the VAE). The fast Mac workarounds are fp16 on MPS (roughly halves memory) or `--max-resolution` to cap the long side; neither is wired as the default. The native-vs-downscale decision lives in the pure helper `invisible_engine._target_size(w, h, max_resolution)` (returns `None` for native, a clamped target tuple otherwise) so it is unit-tested (`tests/test_invisible_engine.py::TestTargetSize`, the #10/#15 regression guard) without loading the model -- keep that logic in the helper, don't re-inline it. +- `invisible` pipeline processes at **native resolution by default** (`max_resolution=0`), matching the hosted raiw.cc backend (fal fast-sdxl, no pre-downscale). The old forced downscale-to-1024 -> upscale-back round-trip was the main quality loss (issue #10) and is gone; at strength ~0.05 SDXL img2img does not need the ~1024 downscale. `--max-resolution N` re-introduces an opt-in long-side cap purely to bound GPU/MPS memory on very large inputs (it reintroduces the lossy round-trip). For huge images that OOM at native, tile-based diffusion is still the proper long-term fix. **Concrete MPS data points (the OOM is memory-tier-dependent, NOT a hard MPS limit):** on a ~24 GB unified-memory machine (verified 2026-05-25, 1254x1254 gpt-image SDXL, fp32) native res OOMs at the *UNet* step (peak ~17 GiB), not only the VAE decode, and the auto-fallback in `img2img_runner` reloads on CPU and finishes (slow, ~13 min) -- the output is still weight-identical and defeats SynthID, so "looks hung/crashed" on Mac is usually this CPU fallback, not a pipeline error. On a **32 GB** unified-memory machine the same default SDXL pass runs entirely on MPS with **no CPU fallback** (verified 2026-05-31, 1122x1402 gpt-image, `all`/default, ~155 s end-to-end), so 32 GB clears the native-res UNet peak that 24 GB could not. Adding `enable_vae_tiling()` alone does NOT prevent the 24 GB OOM (the peak is the UNet, not the VAE). The fast Mac workarounds for memory-constrained machines are fp16 on MPS (roughly halves memory) or `--max-resolution` to cap the long side; neither is wired as the default. **ctrlregen is compute-bound, not memory-bound, on MPS:** the clean-noise profile tiles to 512px (e.g. ~12 tiles for a 1122x1402 image) and runs the full step count per tile (strength 1.0 -> ~50 effective steps/tile), so on a base-tier Apple-GPU laptop it is ~25-30 min/image after the one-time DINOv2-giant download, regardless of having 32 GB -- a discrete CUDA GPU (fp16, no tiling) is the right place for ctrlregen, while the default SDXL pass is comfortable on a 32 GB Mac. The native-vs-downscale decision lives in the pure helper `invisible_engine._target_size(w, h, max_resolution)` (returns `None` for native, a clamped target tuple otherwise) so it is unit-tested (`tests/test_invisible_engine.py::TestTargetSize`, the #10/#15 regression guard) without loading the model -- keep that logic in the helper, don't re-inline it. - **fp16 VAE black-output fix (issue #29, 2026-05-30):** on a **CUDA/XPU fp16** backend the stock SDXL VAE overflows to NaN and the *plain* img2img path decodes to an **all-black** image (reproduced on the raiw.cc result: a 1086x1448 input -> a uniformly black 4.6 KB PNG, mean 0). `watermark_remover._load_pipeline` now swaps in the fp16-fixed SDXL VAE (`madebyollin/sdxl-vae-fp16-fix` = `_SDXL_FP16_VAE_ID`) when `_needs_fp16_vae_fix(model_id, DEFAULT_MODEL_ID, is_fp16)` is true -- only the default SDXL checkpoint on fp16. **cpu/mps run fp32** (the stock VAE is fine there, which is why the bug never reproduces on Mac), and the **differential / region-hires** pipeline already upcasts the VAE itself (see the `text_protector` bullet). A custom non-SDXL `model_id` keeps its own VAE (the fp16-fix VAE is SDXL-architecture-specific). The decision is a pure helper, unit-tested without a download (`tests/test_platform.py::TestFp16VaeFix`); the actual black->clean recovery needs a CUDA GPU and was NOT verifiable on this MPS machine -- confirm on the backend / an NVIDIA box. - Pyright first run is slow (2-3 min) due to ML deps (torch/diffusers/transformers stubs); full-project `uv run pyright` can stall for many minutes — scope it to changed files. - `ultralytics` monkey-patches `PIL.Image.open` and tries to autoload `pi_heif`. When `pi_heif` is missing, opening files raises `ModuleNotFoundError`, not `UnidentifiedImageError`. Code that opens user-supplied or unknown-format files should `except Exception`, not just `OSError`/`UnidentifiedImageError`. diff --git a/pyproject.toml b/pyproject.toml index bc56280..02bcedc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,6 @@ dependencies = [ "numpy>=1.24.0", "opencv-python-headless>=4.8.0", "click>=8.0.0", - "rich>=13.0.0", "python-dotenv>=1.0.0", ] diff --git a/src/remove_ai_watermarks/__init__.py b/src/remove_ai_watermarks/__init__.py index 0673b90..a89d6b7 100644 --- a/src/remove_ai_watermarks/__init__.py +++ b/src/remove_ai_watermarks/__init__.py @@ -1,3 +1,14 @@ """Remove-AI-Watermarks: Unified tool for removing visible and invisible AI watermarks.""" +import os as _os +import warnings as _warnings + +# transformers prints a noisy deprecation for the Siglip2ImageProcessorFast +# alias when it is imported (by the optional GPU/ML path). Silence it before +# any submodule pulls transformers in, so the CLI startup stays quiet. Uses +# setdefault so a user-set TRANSFORMERS_VERBOSITY still wins. +_os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error") +_warnings.filterwarnings("ignore", message=r".*ImageProcessorFast.*") + + __version__ = "0.8.1" diff --git a/src/remove_ai_watermarks/cli.py b/src/remove_ai_watermarks/cli.py index 3652f79..0aadaa6 100644 --- a/src/remove_ai_watermarks/cli.py +++ b/src/remove_ai_watermarks/cli.py @@ -1,13 +1,14 @@ """Unified CLI for remove-ai-watermarks. Provides commands for: - - Visible watermark removal (Gemini sparkle) — works offline, fast - - Invisible watermark removal (SynthID etc.) — requires GPU/diffusion models - - AI metadata stripping — lightweight, no ML deps needed + - Visible watermark removal (Gemini sparkle) - works offline, fast + - Invisible watermark removal (SynthID etc.) - requires GPU/diffusion models + - AI metadata stripping - lightweight, no ML deps needed """ from __future__ import annotations +import contextlib import json import logging import time @@ -15,20 +16,92 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Literal import click -from rich.console import Console -from rich.panel import Panel -from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn, TimeElapsedColumn -from rich.table import Table from remove_ai_watermarks import __version__, watermark_registry -from remove_ai_watermarks.noai.watermark_profiles import DEFAULT_STRENGTH +from remove_ai_watermarks.noai.watermark_profiles import resolve_strength if TYPE_CHECKING: + from collections.abc import Generator + from numpy.typing import NDArray from remove_ai_watermarks.gemini_engine import DetectionResult -console = Console() +# --- plain-text output layer (replaces rich: no colors, no markup, no boxes) --- + + +class _Table: + """Plain-text stand-in for rich.Table.""" + + def __init__(self, *args: Any, title: str | None = None, **kwargs: Any) -> None: + self._title = title + self._headers: list[str] = [] + self._rows: list[list[str]] = [] + + def add_column(self, header: str = "", *args: Any, **kwargs: Any) -> None: + self._headers.append(str(header)) + + def add_row(self, *cells: Any) -> None: + self._rows.append([str(c) for c in cells]) + + def render(self) -> str: + lines: list[str] = [] + if self._title: + lines.append(self._title) + if any(self._headers): + lines.append(" ".join(self._headers)) + lines.extend(" ".join(row) for row in self._rows) + return "\n".join(f" {line}" for line in lines) + + +class _Progress: + """No-op stand-in for rich.Progress; results are printed directly instead.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + pass + + def __enter__(self) -> _Progress: + return self + + def __exit__(self, *exc: object) -> bool: + return False + + def add_task(self, *args: Any, **kwargs: Any) -> int: + return 0 + + def advance(self, *args: Any, **kwargs: Any) -> None: + pass + + def update(self, *args: Any, **kwargs: Any) -> None: + pass + + +class _Console: + """Minimal plain-text replacement for rich.Console.""" + + def print(self, *objects: Any, **kwargs: Any) -> None: + click.echo(" ".join(o.render() if isinstance(o, _Table) else str(o) for o in objects)) + + @contextlib.contextmanager + def status(self, message: str = "", **kwargs: Any) -> Generator[None, None, None]: + if message: + click.echo(message) + yield + + +def _panel(text: str = "", *args: Any, **kwargs: Any) -> str: + return text + + +def _column(*args: Any, **kwargs: Any) -> None: + return None + + +Panel = _panel +Table = _Table +Progress = _Progress +SpinnerColumn = BarColumn = TextColumn = TimeElapsedColumn = _column +console = _Console() SUPPORTED_FORMATS = {".png", ".jpg", ".jpeg", ".webp"} @@ -45,7 +118,7 @@ def _setup_logging(verbose: bool) -> None: def _banner() -> None: console.print( Panel( - f"[bold cyan]Remove-AI-Watermarks[/] [dim]v{__version__}[/]\n[dim]Visible & invisible watermark removal[/]", + f"Remove-AI-Watermarks v{__version__}\nVisible & invisible watermark removal", border_style="cyan", padding=(0, 2), ) @@ -54,12 +127,10 @@ def _banner() -> None: def _validate_image(path: Path) -> Path: if not path.exists(): - console.print(f"[red]Error:[/] File not found: {path}") + console.print(f"Error: File not found: {path}") raise SystemExit(1) if path.suffix.lower() not in SUPPORTED_FORMATS: - console.print( - f"[yellow]Warning:[/] {path.suffix} may not be supported (expected: {', '.join(SUPPORTED_FORMATS)})" - ) + console.print(f"Warning: {path.suffix} may not be supported (expected: {', '.join(SUPPORTED_FORMATS)})") return path @@ -124,7 +195,7 @@ def _write_bgr_with_alpha( image_io.imwrite(path, bgra) -# ── Main group ─────────────────────────────────────────────────────── +# -- Main group ------------------------------------------------------- @click.group(invoke_without_command=True) @@ -146,7 +217,7 @@ def main(ctx: click.Context, verbose: bool) -> None: click.echo(ctx.get_help()) -# ── Visible (Gemini) watermark removal ─────────────────────────────── +# -- Visible (Gemini) watermark removal ------------------------------- @main.command("visible") @@ -198,11 +269,11 @@ def cmd_visible( # Load image (preserving any alpha channel separately) image, alpha = _read_bgr_and_alpha(source) if image is None: - console.print(f"[red]Error:[/] Failed to read image: {source}") + console.print(f"Error: Failed to read image: {source}") raise SystemExit(1) h, w = image.shape[:2] - console.print(f" [dim]Input:[/] {source.name} ({w}x{h})") + console.print(f" Input: {source.name} ({w}x{h})") # Resolve the target mark from the known-watermark registry. ``auto`` scans # every in-auto mark in its usual place and picks the strongest; an explicit @@ -210,31 +281,28 @@ def cmd_visible( if mark == "auto": best = registry.best_auto_mark(image) if best is None: - console.print(" [yellow]⚠[/] No known visible mark detected (gemini / doubao).") + console.print(" Warning: No known visible mark detected (gemini / doubao).") if detect: - console.print(" [dim]Skipping. Use --mark --no-detect to force.[/]") + console.print(" Skipping. Use --mark --no-detect to force.") raise SystemExit(0) target = "gemini" # forced (no-detect): fall back to the default mark else: target = best.key - console.print(f" [dim]Mark auto:[/] {best.label} [dim]({best.location}, conf {best.confidence:.2f})[/]") + console.print(f" Mark auto: {best.label} ({best.location}, conf {best.confidence:.2f})") else: target = mark chosen = registry.get_mark(target) det = chosen.detect(image) if detect and not det.detected: - console.print( - f" [yellow]⚠[/] {chosen.label} not detected " - f"[dim](conf {det.confidence:.2f}). Use --no-detect to force.[/]" - ) + console.print(f" Warning: {chosen.label} not detected (conf {det.confidence:.2f}). Use --no-detect to force.") raise SystemExit(0) if det.detected: - console.print(f" [green]✓[/] {chosen.label} detected [dim]({chosen.location}, conf {det.confidence:.2f})[/]") + console.print(f" {chosen.label} detected ({chosen.location}, conf {det.confidence:.2f})") method: Literal["telea", "ns"] = "ns" if inpaint_method == "ns" else "telea" t0 = time.monotonic() - with console.status(f"[cyan]Removing {chosen.label}… ({chosen.recovery})[/]"): + with console.status(f"Removing {chosen.label}... ({chosen.recovery})"): result, _ = chosen.remove( image, inpaint_method=method, @@ -256,13 +324,13 @@ def cmd_visible( remove_ai_metadata(output, output) except Exception as e: if ctx.obj.get("verbose"): - console.print(f" [yellow]⚠[/] Failed to strip metadata: {e}") + console.print(f" Warning: Failed to strip metadata: {e}") size_kb = output.stat().st_size / 1024 - console.print(f" [green]✓[/] Saved: {output} [dim]({size_kb:.0f} KB, {elapsed:.2f}s)[/]") + console.print(f" Saved: {output} ({size_kb:.0f} KB, {elapsed:.2f}s)") -# ── Universal region eraser ───────────────────────────────────────── +# -- Universal region eraser ----------------------------------------- def _parse_region(spec: str) -> tuple[int, int, int, int]: @@ -322,18 +390,18 @@ def cmd_erase( image, alpha = _read_bgr_and_alpha(source) if image is None: - console.print(f"[red]Error:[/] Failed to read image: {source}") + console.print(f"Error: Failed to read image: {source}") raise SystemExit(1) h, w = image.shape[:2] - console.print(f" [dim]Input:[/] {source.name} ({w}x{h}) [dim]{len(boxes)} region(s), backend={backend}[/]") + console.print(f" Input: {source.name} ({w}x{h}) {len(boxes)} region(s), backend={backend}") t0 = time.monotonic() method: Literal["telea", "ns"] = "ns" if inpaint_method == "ns" else "telea" try: - with console.status(f"[cyan]Erasing ({backend})…[/]"): + with console.status(f"Erasing ({backend})..."): result = erase(image, boxes=boxes, backend=backend, dilate=dilate, cv2_method=method) except RuntimeError as e: - console.print(f" [red]Error:[/] {e}") + console.print(f" Error: {e}") raise SystemExit(1) from e elapsed = time.monotonic() - t0 @@ -347,13 +415,13 @@ def cmd_erase( remove_ai_metadata(output, output) except Exception as e: if ctx.obj.get("verbose"): - console.print(f" [yellow]⚠[/] Failed to strip metadata: {e}") + console.print(f" Warning: Failed to strip metadata: {e}") size_kb = output.stat().st_size / 1024 - console.print(f" [green]✓[/] Erased {len(boxes)} region(s) → {output} [dim]({size_kb:.0f} KB, {elapsed:.2f}s)[/]") + console.print(f" Erased {len(boxes)} region(s) -> {output} ({size_kb:.0f} KB, {elapsed:.2f}s)") -# ── Invisible watermark removal ───────────────────────────────────── +# -- Invisible watermark removal ------------------------------------- @main.command("invisible") @@ -361,7 +429,12 @@ def cmd_erase( @click.option( "-o", "--output", type=click.Path(path_type=Path), default=None, help="Output path (default: _clean.)." ) -@click.option("--strength", type=float, default=DEFAULT_STRENGTH, help="Denoising strength (0.0-1.0). Default: 0.10.") +@click.option( + "--strength", + type=float, + default=None, + help="Denoising strength (0.0-1.0). Default: 0.10 (SDXL), 1.0 clean-noise for ctrlregen.", +) @click.option("--steps", type=int, default=50, help="Number of denoising steps. Default: 50.") @click.option( "--pipeline", @@ -397,7 +470,7 @@ def cmd_invisible( ctx: click.Context, source: Path, output: Path | None, - strength: float, + strength: float | None, steps: int, pipeline: str, device: str, @@ -416,8 +489,7 @@ def cmd_invisible( if not invisible_available(): console.print( - "[red]Error:[/] GPU dependencies not installed.\n" - " Install them with: [bold]pip install 'remove-ai-watermarks\\[gpu]'[/]" + "Error: GPU dependencies not installed.\n Install them with: pip install 'remove-ai-watermarks[gpu]'" ) raise SystemExit(1) @@ -430,7 +502,7 @@ def cmd_invisible( device_str = None if device == "auto" else device def progress_cb(msg: str) -> None: - console.print(f" [dim]{msg}[/]") + console.print(f" {msg}") engine = InvisibleEngine( device=device_str, @@ -439,9 +511,9 @@ def cmd_invisible( progress_callback=progress_cb, ) - console.print(f" [dim]Input:[/] {source.name}") - console.print(f" [dim]Pipeline:[/] {pipeline}") - console.print(f" [dim]Strength:[/] {strength} Steps: {steps}") + console.print(f" Input: {source.name}") + console.print(f" Pipeline: {pipeline}") + console.print(f" Strength: {resolve_strength(strength, pipeline)} Steps: {steps}") t0 = time.monotonic() result_path = engine.remove_watermark( @@ -458,10 +530,10 @@ def cmd_invisible( elapsed = time.monotonic() - t0 size_kb = result_path.stat().st_size / 1024 - console.print(f"\n [green]✓[/] Saved: {result_path} [dim]({size_kb:.0f} KB, {elapsed:.1f}s)[/]") + console.print(f"\n Saved: {result_path} ({size_kb:.0f} KB, {elapsed:.1f}s)") -# ── Metadata operations ───────────────────────────────────────────── +# -- Metadata operations --------------------------------------------- @main.command("metadata") @@ -499,10 +571,10 @@ def cmd_metadata( if check or (not remove): has_ai = has_ai_metadata(source) if has_ai: - console.print(f" [yellow]⚠[/] AI metadata detected in {source.name}:") + console.print(f" Warning: AI metadata detected in {source.name}:") meta = get_ai_metadata(source) if synthid := meta.get("synthid_watermark"): - console.print(f" [bold yellow]⚠ SynthID watermark (inferred from C2PA metadata) {synthid}[/]") + console.print(f" Warning: SynthID watermark (inferred from C2PA metadata) {synthid}") table = Table(show_header=True, header_style="bold") table.add_column("Key", style="cyan") table.add_column("Value") @@ -510,17 +582,17 @@ def cmd_metadata( table.add_row(k, str(v)[:80]) console.print(table) else: - console.print(f" [green]✓[/] No AI metadata found in {source.name}") + console.print(f" No AI metadata found in {source.name}") if not remove: return # Remove out = remove_ai_metadata(source, output, keep_standard=keep_standard) - console.print(f" [green]✓[/] AI metadata stripped → {out}") + console.print(f" AI metadata stripped -> {out}") -# ── Provenance identification ─────────────────────────────────────── +# -- Provenance identification --------------------------------------- @main.command("identify") @@ -552,24 +624,22 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo return _banner() - verdict = {True: "[yellow]AI-generated[/]", False: "[green]not AI[/]", None: "[dim]unknown[/]"}[ - report.is_ai_generated - ] - console.print(f"\n Verdict: {verdict} [dim](confidence: {report.confidence})[/]") - console.print(f" Platform: {report.platform or '[dim]undetermined[/]'}") + verdict = {True: "AI-generated", False: "not AI", None: "unknown"}[report.is_ai_generated] + console.print(f"\n Verdict: {verdict} (confidence: {report.confidence})") + console.print(f" Platform: {report.platform or 'undetermined'}") if report.is_ai_generated is None: console.print( - " [dim]No locally-readable AI signal found. This is not the same as 'clean': " + " No locally-readable AI signal found. This is not the same as 'clean': " "metadata is often stripped by re-encoding, screenshots, or upload, and SynthID-class " "pixel watermarks (Gemini / Nano Banana / gpt-image) have no local detector. " - "See caveats below.[/]" + "See caveats below." ) if report.integrity_clashes: - console.print("\n [bold red]⚠ Integrity clash[/] [dim](provenance signals contradict each other)[/]") + console.print("\n Warning: Integrity clash (provenance signals contradict each other)") for clash in report.integrity_clashes: - console.print(f" [red]- {clash}[/]") + console.print(f" - {clash}") if report.watermarks: table = Table(show_header=True, header_style="bold", title="Watermarks / provenance markers") @@ -578,15 +648,15 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo table.add_row(wm) console.print(table) else: - console.print(" [dim]No watermarks or provenance markers found.[/]") + console.print(" No watermarks or provenance markers found.") if report.caveats: - console.print("\n [dim]Caveats:[/]") + console.print("\n Caveats:") for c in report.caveats: - console.print(f" [dim]- {c}[/]") + console.print(f" - {c}") -# ── Combined "all" mode ────────────────────────────────────────────── +# -- Combined "all" mode ---------------------------------------------- @main.command("all") @@ -599,7 +669,10 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo "--inpaint-method", type=click.Choice(["ns", "telea", "gaussian"]), default="ns", help="Inpainting method." ) @click.option( - "--strength", type=float, default=DEFAULT_STRENGTH, help="Invisible watermark denoising strength. Default: 0.10." + "--strength", + type=float, + default=None, + help="Invisible watermark denoising strength. Default: 0.10 (SDXL), 1.0 clean-noise for ctrlregen.", ) @click.option("--steps", type=int, default=50, help="Number of denoising steps for invisible removal.") @click.option( @@ -639,7 +712,7 @@ def cmd_all( output: Path | None, inpaint: bool, inpaint_method: Literal["ns", "telea", "gaussian"], - strength: float, + strength: float | None, steps: int, pipeline: str, model: str | None, @@ -680,40 +753,40 @@ def cmd_all( os.close(tmp_fd) - # ── Step 1: Visible watermark ──────────────────────────────── - console.print("\n [bold cyan]① Visible watermark removal[/]") + # -- Step 1: Visible watermark -------------------------------- + console.print("\n 1) Visible watermark removal") engine = GeminiEngine() image, alpha = _read_bgr_and_alpha(source) if image is None: - console.print(f"[red]Error:[/] Failed to read image: {source}") + console.print(f"Error: Failed to read image: {source}") raise SystemExit(1) h, w = image.shape[:2] - console.print(f" [dim]Input:[/] {source.name} ({w}x{h})") + console.print(f" Input: {source.name} ({w}x{h})") - with console.status("[cyan]Removing visible watermark…[/]"): + with console.status("Removing visible watermark..."): det = engine.detect_watermark(image) if det.detected: result = engine.remove_watermark(image) if inpaint: region = _watermark_region(det, w, h) result = engine.inpaint_residual(result, region, method=inpaint_method) - console.print(" [green]✓[/] Visible watermark removed") + console.print(" Visible watermark removed") else: result = image.copy() - console.print(" [dim]Skipped (no visible watermark detected)[/]") + console.print(" Skipped (no visible watermark detected)") # Save to temp file for invisible engine input (preserve alpha if present) _write_bgr_with_alpha(tmp_path, result, alpha) - # ── Step 2: Invisible watermark ────────────────────────────── - console.print("\n [bold cyan]② Invisible watermark removal[/]") + # -- Step 2: Invisible watermark ------------------------------ + console.print("\n 2) Invisible watermark removal") from remove_ai_watermarks.invisible_engine import is_available as invisible_available if not invisible_available(): console.print( - " [yellow]⚠[/] Skipped — GPU dependencies not installed.\n" - " Install them with: [bold]pip install 'remove-ai-watermarks\\[gpu]'[/]" + " Warning: Skipped - GPU dependencies not installed.\n" + " Install them with: pip install 'remove-ai-watermarks[gpu]'" ) else: from remove_ai_watermarks.invisible_engine import InvisibleEngine @@ -721,7 +794,7 @@ def cmd_all( device_str = None if device == "auto" else device def progress_cb(msg: str) -> None: - console.print(f" [dim]{msg}[/]") + console.print(f" {msg}") inv_engine = InvisibleEngine( model_id=model, @@ -731,7 +804,7 @@ def cmd_all( progress_callback=progress_cb, ) - console.print(f" [dim]Strength:[/] {strength} Steps: {steps}") + console.print(f" Strength: {resolve_strength(strength, pipeline)} Steps: {steps}") inv_engine.remove_watermark( image_path=tmp_path, output_path=tmp_path, @@ -742,26 +815,26 @@ def cmd_all( protect_text=not no_protect_text, max_resolution=max_resolution, ) - console.print(" [green]✓[/] Invisible watermark removed") + console.print(" Invisible watermark removed") - # ── Step 3: Metadata ───────────────────────────────────────── - console.print("\n [bold cyan]③ AI metadata stripping[/]") + # -- Step 3: Metadata ----------------------------------------- + console.print("\n 3) AI metadata stripping") try: from remove_ai_watermarks.metadata import remove_ai_metadata remove_ai_metadata(tmp_path, tmp_path) - console.print(" [green]✓[/] AI metadata stripped") + console.print(" AI metadata stripped") except Exception as e: - console.print(f" [yellow]⚠[/] Metadata strip failed: {e}") + console.print(f" Warning: Metadata strip failed: {e}") - # ── Write final result ──────────────────────────────────────── + # -- Write final result ---------------------------------------- # The invisible step (and downstream cv2.IMREAD_COLOR paths) drops alpha, # so re-attach the original alpha plane unchanged when writing the final # output for transparent formats. output.parent.mkdir(parents=True, exist_ok=True) final_bgr, _ = _read_bgr_and_alpha(tmp_path) if final_bgr is None: - console.print(f"[red]Error:[/] Failed to read intermediate file: {tmp_path}") + console.print(f"Error: Failed to read intermediate file: {tmp_path}") raise SystemExit(1) _write_bgr_with_alpha(output, final_bgr, alpha) @@ -770,13 +843,13 @@ def cmd_all( if tmp_path.exists(): tmp_path.unlink() - # ── Done ───────────────────────────────────────────────────── + # -- Done ----------------------------------------------------- elapsed = time.monotonic() - t0 size_kb = output.stat().st_size / 1024 - console.print(f"\n [bold green]✓ Done:[/] {output} [dim]({size_kb:.0f} KB, {elapsed:.1f}s total)[/]") + console.print(f"\n Done: {output} ({size_kb:.0f} KB, {elapsed:.1f}s total)") -# ── Batch command ──────────────────────────────────────────────────── +# -- Batch command ---------------------------------------------------- def _process_batch_image( @@ -932,12 +1005,12 @@ def cmd_batch( images = sorted(p for p in directory.iterdir() if p.suffix.lower() in SUPPORTED_FORMATS) if not images: - console.print(f"[yellow]No supported images found in {directory}[/]") + console.print(f"No supported images found in {directory}") return - console.print(f" Found [bold]{len(images)}[/] images in {directory}") - console.print(f" Output → {output_dir}") - console.print(f" Mode: [cyan]{mode}[/]") + console.print(f" Found {len(images)} images in {directory}") + console.print(f" Output -> {output_dir}") + console.print(f" Mode: {mode}") processed = 0 errors = 0 @@ -950,11 +1023,11 @@ def cmd_batch( TimeElapsedColumn(), console=console, ) as progress: - task = progress.add_task("Processing…", total=len(images)) + task = progress.add_task("Processing...", total=len(images)) for img_path in images: out_path = output_dir / img_path.name - progress.update(task, description=f"[cyan]{img_path.name}[/]") + progress.update(task, description=f"{img_path.name}") try: _process_batch_image( @@ -977,11 +1050,11 @@ def cmd_batch( except Exception as e: errors += 1 if ctx.obj.get("verbose"): - console.print(f" [red]✗[/] {img_path.name}: {e}") + console.print(f" {img_path.name}: {e}") progress.advance(task) - console.print(f"\n [green]✓[/] {processed} processed" + (f" [red]✗[/] {errors} errors" if errors else "")) + console.print(f"\n {processed} processed" + (f" {errors} errors" if errors else "")) if __name__ == "__main__": diff --git a/tests/test_cli.py b/tests/test_cli.py index 98dc520..de922e4 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -526,8 +526,8 @@ class TestBatchCommand: class TestGpuHintMarkup: - """The GPU-extra install hint must survive rich markup (the ``[gpu]`` token - is otherwise parsed as a style tag and silently dropped).""" + """The GPU-extra install hint must reach the user with the ``[gpu]`` token + intact (plain output prints it verbatim, with no markup parsing).""" def test_invisible_install_hint_keeps_gpu_extra(self, runner, sample_png): with patch("remove_ai_watermarks.invisible_engine.is_available", return_value=False): diff --git a/uv.lock b/uv.lock index f70ebd7..72d05da 100644 --- a/uv.lock +++ b/uv.lock @@ -2874,7 +2874,6 @@ dependencies = [ { name = "piexif" }, { name = "pillow" }, { name = "python-dotenv" }, - { name = "rich" }, ] [package.optional-dependencies] @@ -2948,7 +2947,6 @@ requires-dist = [ { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.1.0" }, { name = "python-dotenv", specifier = ">=1.0.0" }, { name = "remove-ai-watermarks", extras = ["gpu", "detect", "trustmark", "lama", "dev"], marker = "extra == 'all'" }, - { name = "rich", specifier = ">=13.0.0" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.4.0" }, { name = "safetensors", marker = "extra == 'gpu'" }, { name = "tokenizers", marker = "extra == 'gpu'", specifier = ">=0.22,<0.23" },