remove-ai-watermarks/src/remove_ai_watermarks/_internal/progress.py

"""Progress reporting utilities for long-running optional model operations."""

from __future__ import annotations

import contextlib
import io
import os
import sys
import threading
import time
import warnings
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from collections.abc import Callable

_BAR_WIDTH = 28
_SPINNER = ("|", "/", "-", "\\")


def _truncate(text: str, max_len: int = 72) -> str:
    if len(text) <= max_len:
        return text
    return f"{text[: max(0, max_len - 3)]}..."


def _build_bar(step: int) -> str:
    position = step % (2 * _BAR_WIDTH - 2)
    if position >= _BAR_WIDTH:
        position = 2 * _BAR_WIDTH - 2 - position
    cells = ["-"] * _BAR_WIDTH
    cells[position] = "="
    return "".join(cells)


@dataclass
class _TaskResult:
    value: Any = None
    error: BaseException | None = None
    complete: threading.Event = field(default_factory=threading.Event)


def run_with_progress(task: Callable[[], Any], progress_state: dict[str, str] | None = None) -> Any:
    """Run ``task`` on a worker thread and render a compact terminal heartbeat."""
    outcome = _TaskResult()

    def invoke() -> None:
        try:
            outcome.value = task()
        except BaseException as error:  # re-raised on the caller thread
            outcome.error = error
        finally:
            outcome.complete.set()

    worker = threading.Thread(target=invoke, name="raiw-progress-task", daemon=True)
    worker.start()
    started_at = time.monotonic()
    frame = 0
    terminal = sys.__stderr__
    while not outcome.complete.wait(0.1):
        message = _truncate((progress_state or {}).get("message", "Processing..."))
        elapsed = int(time.monotonic() - started_at)
        if terminal is not None:
            terminal.write(
                f"\r\033[2K  {_SPINNER[frame % len(_SPINNER)]} [{_build_bar(frame)}] {elapsed:>3}s  {message}"
            )
            terminal.flush()
        frame += 1

    worker.join()
    elapsed = int(time.monotonic() - started_at)
    message = _truncate((progress_state or {}).get("message", "Processing..."))
    if terminal is not None:
        terminal.write(f"\r\033[2K  Completed in {elapsed}s  {message}\n")
        terminal.flush()
    if outcome.error is not None:
        raise outcome.error
    return outcome.value


def _silence_diffusers() -> None:
    from diffusers.utils import logging as diffusers_logging

    diffusers_logging.set_verbosity_error()
    disable = getattr(diffusers_logging, "disable_progress_bar", None)
    if callable(disable):
        disable()


def _configure_quiet_libraries() -> None:
    os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
    operations = (
        lambda: __import__("transformers").logging.set_verbosity_error(),
        _silence_diffusers,
        lambda: __import__("huggingface_hub").logging.set_verbosity_error(),
    )
    for operation in operations:
        with contextlib.suppress(Exception):
            operation()


def silence_library_output(
    run_func: Callable[[], Any],
    set_progress: Callable[[str], None] | None = None,
) -> Callable[[], Any]:
    """Wrap a model call so third-party progress bars do not corrupt our CLI UI."""

    def quiet_call() -> Any:
        if set_progress is not None:
            set_progress("Preparing model runtime...")
        _configure_quiet_libraries()
        with (
            warnings.catch_warnings(),
            contextlib.redirect_stdout(io.StringIO()),
            contextlib.redirect_stderr(io.StringIO()),
        ):
            warnings.simplefilter("ignore")
            if set_progress is not None:
                set_progress("Running watermark regeneration...")
            return run_func()

    return quiet_call


@dataclass
class _PipelineMonitor:
    total_steps: int
    device: str
    update: Callable[[str], None]
    bar_len: int
    label: str
    pre_phases: list[tuple[int, str]]
    post_phases: list[tuple[int, str]]
    first_step: threading.Event = field(default_factory=threading.Event)
    done: threading.Event = field(default_factory=threading.Event)
    started_at: float = field(default_factory=time.monotonic)
    last_step_at: float = field(default_factory=time.monotonic)

    def callback(self, step: int, _timestep: int, _latents: Any) -> None:
        self.first_step.set()
        now = time.monotonic()
        self.last_step_at = now
        current = min(self.total_steps, step + 1)
        filled = round(self.bar_len * current / self.total_steps)
        elapsed = now - self.started_at
        eta = elapsed * max(0, self.total_steps - current) / max(1, current)
        bar = "#" * filled + "." * (self.bar_len - filled)
        self.update(
            f"{self.label} [{bar}] {current}/{self.total_steps}, "
            f"{elapsed:.0f}s elapsed, ~{eta:.0f}s left, {self.device}"
        )

    def _phase_message(self, phases: list[tuple[int, str]], elapsed: float) -> str:
        message = phases[0][1]
        for threshold, candidate in phases:
            if elapsed < threshold:
                break
            message = candidate
        return message

    def monitor(self) -> None:
        while not self.first_step.wait(0.4):
            elapsed = time.monotonic() - self.started_at
            self.update(self._phase_message(self.pre_phases, elapsed))
        decode_started: float | None = None
        while not self.done.wait(0.4):
            if time.monotonic() - self.last_step_at < 1.5:
                decode_started = None
                continue
            decode_started = decode_started or time.monotonic()
            self.update(self._phase_message(self.post_phases, time.monotonic() - decode_started))

    def start(self) -> threading.Thread:
        self.started_at = self.last_step_at = time.monotonic()
        self.first_step.clear()
        self.done.clear()
        thread = threading.Thread(target=self.monitor, name="raiw-pipeline-progress", daemon=True)
        thread.start()
        return thread


def make_pipeline_progress(
    effective_steps: int,
    device: str,
    set_progress: Callable[[str], None],
    *,
    bar_len: int = 20,
    label: str = "Denoising",
    pre_phases: list[tuple[int, str]] | None = None,
    post_phases: list[tuple[int, str]] | None = None,
) -> tuple[Callable[..., None], threading.Event, threading.Event, Callable[[], threading.Thread]]:
    """Build a callback and monitor for the legacy Diffusers callback interface."""

    def qualify(entries: list[tuple[int, str]]) -> list[tuple[int, str]]:
        return [(second, f"{text} on {device}") for second, text in entries]

    monitor = _PipelineMonitor(
        total_steps=max(1, effective_steps),
        device=device,
        update=set_progress,
        bar_len=bar_len,
        label=label,
        pre_phases=pre_phases or qualify([(0, "Encoding image"), (8, "Preparing denoiser"), (20, "Starting sampler")]),
        post_phases=post_phases or qualify([(0, "Decoding image"), (10, "Finalizing pixels"), (45, "Still decoding")]),
    )
    return monitor.callback, monitor.first_step, monitor.done, monitor.start


def is_mps_error(error: Exception) -> bool:
    """Return whether an error message identifies Apple's MPS backend."""
    return "mps" in str(error).casefold()