mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-10 16:10:33 +02:00
Remove the unreachable ESRGAN upscale chain
The min-resolution floor lifted small inputs toward SDXL's ~1024 training size, and Real-ESRGAN was an optional way to do that lifting. Both surviving profiles run at native geometry, so the engine forced the floor to 0 on every path; the floor never fired, `upscaling` was never true, and nothing downstream of it could execute. Gone: upscaler.py, _esrgan_upscale, the min_resolution and upscaler parameters, _target_size's floor branch, --min-resolution, --upscaler, _warn_if_esrgan_unavailable and the `esrgan` extra. max_resolution stays and is now the only lever on geometry; it can only scale down. scripts/smoke_matrix.py was the one live consumer and neither gate saw it - Pyright is scoped to src/ and Ruff cannot resolve its function-local import - so `--diffusion` would have died at import. Its knob rows were written for the removed profiles besides (--pipeline sdxl, --steps 20, --guidance-scale 5.0, --device mps), so they are rewritten rather than patched: most now assert a knob is REJECTED, which is the coverage worth having when the CLI accepts a value the library refuses several layers down. Accepted-knob rows skip without CUDA, so the row count is host-dependent and verification-plan.md no longer claims a fixed 68. This removes a public module, a CLI option and a published extra, so the next release is 0.25.0, not a patch. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
co-authored by
Claude Opus 5
Claude Opus 5
parent
95a6964e04
commit
bf4bfc1ab7
@@ -180,28 +180,10 @@ _controlnet_scale_option = click.option(
|
||||
"(structure/text preservation strength). Higher = closer to original structure.",
|
||||
)
|
||||
|
||||
_min_resolution_option = click.option(
|
||||
"--min-resolution",
|
||||
type=int,
|
||||
default=1024,
|
||||
help="Upscale long side UP to this (px) before diffusion when the input is smaller, so SDXL runs "
|
||||
"near 1024 (small inputs distort at native); output is restored to the input size. 0 = off. Default 1024.",
|
||||
)
|
||||
|
||||
_unsharp_option = click.option(
|
||||
"--unsharp", type=float, default=0.0, help="Unsharp-mask sharpening strength (0 = off, typical: 0.3-0.8)."
|
||||
)
|
||||
|
||||
_upscaler_option = click.option(
|
||||
"--upscaler",
|
||||
type=click.Choice(["lanczos", "esrgan"]),
|
||||
default="lanczos",
|
||||
help="How to upscale a small input to the --min-resolution floor: lanczos (default, cv2, no model) or "
|
||||
"esrgan (Real-ESRGAN via the 'esrgan' extra; better detail, slower on CPU). Best for photo/texture "
|
||||
"content -- as a generic GAN with no face/glyph prior it can degrade faces (diffusion mitigates) and "
|
||||
"thin text, so lanczos stays the default. Falls back to lanczos if the extra is absent. Only when upscaling.",
|
||||
)
|
||||
|
||||
_auto_option = click.option(
|
||||
"--auto",
|
||||
is_flag=True,
|
||||
@@ -376,21 +358,6 @@ def _resolve_profile_polish(auto: bool, adaptive_polish: bool, pipeline: str) ->
|
||||
return adaptive_polish
|
||||
|
||||
|
||||
def _warn_if_esrgan_unavailable(upscaler: str) -> None:
|
||||
"""Tell the user once if ``--upscaler esrgan`` will silently fall back to Lanczos.
|
||||
|
||||
The engine downgrades to Lanczos when the ``esrgan`` extra is absent (fail-safe, so
|
||||
a batch never breaks mid-run) -- but without this notice the user would believe
|
||||
Real-ESRGAN ran. Surfaced at the CLI layer, once per invocation (not per image).
|
||||
"""
|
||||
if upscaler != "esrgan":
|
||||
return
|
||||
from remove_ai_watermarks import upscaler as _upscaler
|
||||
|
||||
if not _upscaler.is_available():
|
||||
console.print(" Note: --upscaler esrgan needs the 'esrgan' extra; falling back to Lanczos.")
|
||||
|
||||
|
||||
def _visible_provenance(path: Path | None) -> frozenset[str]:
|
||||
"""Vendor keys local metadata confirms, the EVIDENCE that drives ``auto``
|
||||
sensitivity. Thin wrapper over the public :func:`api.visible_provenance` (one
|
||||
@@ -896,9 +863,7 @@ def cmd_erase(
|
||||
help="Cap long side (px) before diffusion; 0 = native and preserves the most detail. Raise only on GPU/MPS OOM.",
|
||||
)
|
||||
@_controlnet_scale_option
|
||||
@_min_resolution_option
|
||||
@_unsharp_option
|
||||
@_upscaler_option
|
||||
@_model_option
|
||||
@_guidance_scale_option
|
||||
@_auto_option
|
||||
@@ -920,9 +885,7 @@ def cmd_invisible(
|
||||
humanize: float,
|
||||
unsharp: float,
|
||||
max_resolution: int,
|
||||
min_resolution: int,
|
||||
controlnet_scale: float,
|
||||
upscaler: str,
|
||||
model: str | None,
|
||||
guidance_scale: float | None,
|
||||
auto: bool,
|
||||
@@ -952,7 +915,6 @@ def cmd_invisible(
|
||||
source = _validate_image(source)
|
||||
steps = resolve_steps(steps)
|
||||
seed = resolve_seed(seed)
|
||||
_warn_if_esrgan_unavailable(upscaler)
|
||||
adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline)
|
||||
if output is None:
|
||||
output = source.with_stem(source.stem + "_clean")
|
||||
@@ -998,8 +960,6 @@ def cmd_invisible(
|
||||
unsharp=unsharp,
|
||||
adaptive_polish=adaptive_polish,
|
||||
max_resolution=max_resolution,
|
||||
min_resolution=min_resolution,
|
||||
upscaler=upscaler,
|
||||
vendor=vendor,
|
||||
tile=tile,
|
||||
tile_size=tile_size,
|
||||
@@ -1587,9 +1547,7 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo
|
||||
help="Cap long side (px) before diffusion; 0 = native and preserves the most detail. Raise only on GPU/MPS OOM.",
|
||||
)
|
||||
@_controlnet_scale_option
|
||||
@_min_resolution_option
|
||||
@_unsharp_option
|
||||
@_upscaler_option
|
||||
@_guidance_scale_option
|
||||
@_auto_option
|
||||
@_adaptive_polish_option
|
||||
@@ -1613,9 +1571,7 @@ def cmd_all(
|
||||
humanize: float,
|
||||
unsharp: float,
|
||||
max_resolution: int,
|
||||
min_resolution: int,
|
||||
controlnet_scale: float,
|
||||
upscaler: str,
|
||||
guidance_scale: float | None,
|
||||
auto: bool,
|
||||
adaptive_polish: bool,
|
||||
@@ -1638,7 +1594,6 @@ def cmd_all(
|
||||
source = _validate_image(source)
|
||||
steps = resolve_steps(steps)
|
||||
seed = resolve_seed(seed)
|
||||
_warn_if_esrgan_unavailable(upscaler)
|
||||
adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline)
|
||||
|
||||
if output is None:
|
||||
@@ -1744,8 +1699,6 @@ def cmd_all(
|
||||
unsharp=unsharp,
|
||||
adaptive_polish=adaptive_polish,
|
||||
max_resolution=max_resolution,
|
||||
min_resolution=min_resolution,
|
||||
upscaler=upscaler,
|
||||
vendor=vendor,
|
||||
tile=tile,
|
||||
tile_size=tile_size,
|
||||
@@ -1838,9 +1791,7 @@ class _BatchOptions:
|
||||
sensitivity: str = "auto"
|
||||
unsharp: float = 0.0
|
||||
max_resolution: int = 0
|
||||
min_resolution: int = 1024
|
||||
controlnet_scale: float = 1.0
|
||||
upscaler: str = "lanczos"
|
||||
model: str | None = None
|
||||
guidance_scale: float | None = None
|
||||
adaptive_polish: bool = False
|
||||
@@ -1895,8 +1846,6 @@ def _run_batch_invisible(
|
||||
unsharp=options.unsharp,
|
||||
adaptive_polish=options.adaptive_polish,
|
||||
max_resolution=options.max_resolution,
|
||||
min_resolution=options.min_resolution,
|
||||
upscaler=options.upscaler,
|
||||
tile=options.tile,
|
||||
tile_size=options.tile_size,
|
||||
tile_overlap=options.tile_overlap,
|
||||
@@ -2030,9 +1979,7 @@ def _process_batch_image(
|
||||
default=0,
|
||||
help="Cap long side (px) before diffusion; 0 = native and preserves the most detail. Raise only on GPU/MPS OOM.",
|
||||
)
|
||||
@_min_resolution_option
|
||||
@_unsharp_option
|
||||
@_upscaler_option
|
||||
@_controlnet_scale_option
|
||||
@_model_option
|
||||
@_guidance_scale_option
|
||||
@@ -2058,9 +2005,7 @@ def cmd_batch(
|
||||
humanize: float,
|
||||
unsharp: float,
|
||||
max_resolution: int,
|
||||
min_resolution: int,
|
||||
controlnet_scale: float,
|
||||
upscaler: str,
|
||||
model: str | None,
|
||||
guidance_scale: float | None,
|
||||
auto: bool,
|
||||
@@ -2087,8 +2032,6 @@ def cmd_batch(
|
||||
console.print(f" Found {len(images)} images in {directory}")
|
||||
console.print(f" Output -> {output_dir}")
|
||||
console.print(f" Mode: {mode}")
|
||||
if mode in ("invisible", "all"):
|
||||
_warn_if_esrgan_unavailable(upscaler)
|
||||
adaptive_polish = _resolve_profile_polish(auto, adaptive_polish, pipeline)
|
||||
steps = resolve_steps(steps)
|
||||
seed = resolve_seed(seed)
|
||||
@@ -2104,9 +2047,7 @@ def cmd_batch(
|
||||
sensitivity=sensitivity,
|
||||
unsharp=unsharp,
|
||||
max_resolution=max_resolution,
|
||||
min_resolution=min_resolution,
|
||||
controlnet_scale=controlnet_scale,
|
||||
upscaler=upscaler,
|
||||
model=model,
|
||||
guidance_scale=guidance_scale,
|
||||
adaptive_polish=adaptive_polish,
|
||||
|
||||
@@ -14,7 +14,7 @@ import logging
|
||||
import os
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ._internal.watermark_profiles import (
|
||||
DEFAULT_MODEL_ID as DEFAULT_SDXL_MODEL_ID,
|
||||
@@ -48,21 +48,19 @@ def is_available() -> bool:
|
||||
return module_available("diffusers", "torch")
|
||||
|
||||
|
||||
def _target_size(width: int, height: int, max_resolution: int, min_resolution: int = 0) -> tuple[int, int] | None:
|
||||
def _target_size(width: int, height: int, max_resolution: int) -> tuple[int, int] | None:
|
||||
"""Compute the (width, height) to process at, or None for native.
|
||||
|
||||
Two opposite long-side adjustments, in precedence order:
|
||||
One long-side adjustment: if it exceeds ``max_resolution``, scale DOWN to it
|
||||
(integer-truncated, matching the PIL ``resize`` call site). 0/negative = no cap.
|
||||
Set only to bound GPU/MPS memory on very large inputs (issue #10).
|
||||
|
||||
- ``max_resolution`` (cap): if the long side exceeds it, scale DOWN to it
|
||||
(integer-truncated, matching the PIL ``resize`` call site). 0/negative = no
|
||||
cap. Set only to bound GPU/MPS memory on very large inputs (issue #10).
|
||||
- ``min_resolution`` (floor): else if the long side is below it, scale UP to it
|
||||
(rounded) so SDXL img2img runs near its ~1024 training resolution instead of
|
||||
degrading on a tiny latent (a 381x512 portrait distorts badly at native).
|
||||
The output is restored to the original size by the caller, so the floor is a
|
||||
transparent quality boost. 0 = no floor. Skipped on a ``min > max`` misconfig.
|
||||
There was also a ``min_resolution`` floor that scaled small inputs UP toward
|
||||
SDXL's ~1024 training size. It went with the SDXL profiles: both surviving
|
||||
profiles run at native geometry, so the floor was forced to 0 on every path and
|
||||
could not fire.
|
||||
|
||||
Returns None when neither applies (native resolution). Pure function so the
|
||||
Returns None when the cap does not apply (native resolution). Pure function so the
|
||||
resolution decision is unit-testable without loading the diffusion model.
|
||||
"""
|
||||
long_side = max(width, height)
|
||||
@@ -71,9 +69,6 @@ def _target_size(width: int, height: int, max_resolution: int, min_resolution: i
|
||||
# Clamp the short side to >=1: extreme aspect ratios (e.g. 5000x3 capped
|
||||
# at 1024) would otherwise truncate it to 0 and crash image.resize().
|
||||
return (max(1, int(width * ratio)), max(1, int(height * ratio)))
|
||||
if min_resolution > 0 and long_side < min_resolution and (max_resolution <= 0 or min_resolution <= max_resolution):
|
||||
ratio = min_resolution / long_side
|
||||
return (max(1, round(width * ratio)), max(1, round(height * ratio)))
|
||||
return None
|
||||
|
||||
|
||||
@@ -144,32 +139,6 @@ class InvisibleEngine:
|
||||
"""
|
||||
self._remover.preload(global_only=global_only)
|
||||
|
||||
def _esrgan_upscale(self, image: Any, target: tuple[int, int]) -> Any:
|
||||
"""Upscale a PIL image to ``target`` with Real-ESRGAN, else Lanczos.
|
||||
|
||||
Runs Real-ESRGAN at its native factor (on the remover's device, CPU fallback),
|
||||
then resizes to the exact ``target`` with Lanczos. Falls back to a plain Lanczos
|
||||
resize when the ``esrgan`` extra is absent or the model errors.
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from remove_ai_watermarks import upscaler
|
||||
|
||||
if not upscaler.is_available():
|
||||
logger.debug("esrgan upscaler requested but the extra is absent; using Lanczos")
|
||||
return image.resize(target, Image.Resampling.LANCZOS)
|
||||
try:
|
||||
bgr = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
|
||||
big = upscaler.upscale(bgr, device=self._remover.device)
|
||||
if (big.shape[1], big.shape[0]) != target:
|
||||
big = cv2.resize(big, target, interpolation=cv2.INTER_LANCZOS4)
|
||||
return Image.fromarray(cv2.cvtColor(big, cv2.COLOR_BGR2RGB))
|
||||
except Exception as e: # never let an optional upscaler break removal
|
||||
logger.warning("Real-ESRGAN upscale failed (%s); using Lanczos", e)
|
||||
return image.resize(target, Image.Resampling.LANCZOS)
|
||||
|
||||
def remove_watermark(
|
||||
self,
|
||||
image_path: Path,
|
||||
@@ -180,11 +149,9 @@ class InvisibleEngine:
|
||||
seed: int | None = None,
|
||||
humanize: float = 0.0,
|
||||
max_resolution: int = 0,
|
||||
min_resolution: int = 1024,
|
||||
vendor: str | None = None,
|
||||
unsharp: float = 0.0,
|
||||
adaptive_polish: bool = False,
|
||||
upscaler: str = "lanczos",
|
||||
tile: bool = False,
|
||||
tile_size: int = 1024,
|
||||
tile_overlap: int = 128,
|
||||
@@ -215,17 +182,6 @@ class InvisibleEngine:
|
||||
= no cap. Set a positive value only to bound GPU/MPS memory on
|
||||
very large inputs (it reintroduces a lossy downscale->upscale
|
||||
round-trip).
|
||||
min_resolution: Upscale the long side UP to this (px) before diffusion
|
||||
when the input is smaller, so SDXL runs near its ~1024 training
|
||||
resolution (small inputs degrade/distort badly at native). 1024
|
||||
(default) = on; 0 = off. The output is restored to the original
|
||||
input size, so this is a transparent quality boost; it adds time
|
||||
and memory on small inputs. Ignored on a min > max misconfig.
|
||||
upscaler: How to upscale a small input to the ``min_resolution`` floor:
|
||||
``"lanczos"`` (default, cv2, no model download) or ``"esrgan"`` (Real-ESRGAN
|
||||
via the ``esrgan`` extra). Only applies when UPscaling (the floor
|
||||
case); a ``max_resolution`` downscale always uses Lanczos. Falls back
|
||||
to Lanczos if the extra is absent.
|
||||
tile: Process the diffusion pass in overlapping tiles instead of one
|
||||
forward pass. This retains the input's native dimensions instead
|
||||
of applying ``max_resolution``, but each tile is still regenerated.
|
||||
@@ -244,10 +200,7 @@ class InvisibleEngine:
|
||||
from PIL import Image, ImageOps
|
||||
|
||||
# Resolution policy: a max_resolution cap (0 = none) bounds memory on huge
|
||||
# inputs, and a min_resolution floor (1024 = default) upscales tiny inputs so
|
||||
# SDXL img2img runs near its ~1024 training size instead of distorting on a
|
||||
# tiny latent (a 381x512 portrait wrecks at native -- issue #36 follow-up).
|
||||
# The output is restored to orig_size below, so the floor is transparent.
|
||||
# inputs. See _target_size for why it is the only lever left.
|
||||
# Register the HEIF/AVIF opener so a .heic/.avif input (now a SUPPORTED_FORMAT)
|
||||
# decodes here too. The --force skip path bypasses image_io.imread, which is
|
||||
# what would otherwise register it, so a bare Image.open would fail on HEIC.
|
||||
@@ -261,34 +214,16 @@ class InvisibleEngine:
|
||||
# reassigned to the resized copy below; PIL resize returns a new object).
|
||||
reference_pil = image
|
||||
|
||||
# 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) in {"qwen-zimage", "sdxl-zimage"} else min_resolution
|
||||
)
|
||||
target = _target_size(
|
||||
image.width,
|
||||
image.height,
|
||||
max_resolution,
|
||||
effective_min_resolution,
|
||||
)
|
||||
# Both profiles run at the input's native geometry, so only the explicit max
|
||||
# cap can move it, and it can only ever scale down.
|
||||
target = _target_size(image.width, image.height, max_resolution)
|
||||
if target is not None:
|
||||
upscaling = max(target) > max(image.width, image.height)
|
||||
if self._progress_callback:
|
||||
reason = (
|
||||
f"min-resolution floor {min_resolution}px"
|
||||
if upscaling
|
||||
else f"max-resolution cap {max_resolution}px"
|
||||
self._progress_callback(
|
||||
f"Downscaling {image.width}x{image.height} to {target[0]}x{target[1]} "
|
||||
f"(max-resolution cap {max_resolution}px)..."
|
||||
)
|
||||
verb = "Upscaling" if upscaling else "Downscaling"
|
||||
self._progress_callback(f"{verb} {image.width}x{image.height} to {target[0]}x{target[1]} ({reason})...")
|
||||
# Real-ESRGAN only helps when UPscaling (the floor case); a downscale cap
|
||||
# always uses Lanczos. _esrgan_upscale falls back to Lanczos if the extra is absent.
|
||||
if upscaling and upscaler == "esrgan":
|
||||
image = self._esrgan_upscale(image, target)
|
||||
else:
|
||||
image = image.resize(target, Image.Resampling.LANCZOS)
|
||||
image = image.resize(target, Image.Resampling.LANCZOS)
|
||||
|
||||
# Always persist to a temp file, even without downscaling: WatermarkRemover
|
||||
# reloads by path, so the EXIF-transposed pixels must be saved or rotation
|
||||
|
||||
@@ -1,126 +0,0 @@
|
||||
"""Optional pre-diffusion super-resolution for small inputs (Real-ESRGAN via spandrel).
|
||||
|
||||
Mirrors ``region_eraser``'s optional-backend pattern: ``is_available()`` guards the
|
||||
``spandrel`` import, a lazy singleton (double-checked lock) holds the loaded model, and
|
||||
the weights download on first use (cached by ``torch.hub``) -- they are never bundled.
|
||||
|
||||
The DEFAULT upscaler stays Lanczos (cv2, no model download); this is opt-in via the ``esrgan``
|
||||
extra and feeds the ``--upscaler esrgan`` path. ``spandrel`` is a pure model-loader
|
||||
(MIT) with NO basicsr dependency -- it pulls only torch/torchvision/safetensors/numpy/
|
||||
einops -- so it sidesteps the basicsr / ``torchvision.transforms.functional_tensor``
|
||||
breakage that the retired ``restore`` (GFPGAN) extra had to shim. Real-ESRGAN weights
|
||||
are BSD-3-Clause.
|
||||
|
||||
CPU works but is slow on large inputs, so this is meant for the pre-diffusion upscale of
|
||||
SMALL inputs (and the GPU worker). On a memory-constrained host it is a no-op (the extra
|
||||
is absent), and the caller falls back to Lanczos.
|
||||
"""
|
||||
|
||||
# torch/spandrel boundary: these libs ship no usable element types; relax the
|
||||
# unknown-type rules for this file only.
|
||||
# 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, reportAttributeAccessIssue=false, reportPrivateImportUsage=false
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from numpy.typing import NDArray
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Real-ESRGAN x2plus (BSD-3-Clause), official release. x2 is the right native factor for
|
||||
# the pre-diffusion floor upscale (small inputs ~512 -> ~1024); spandrel infers the
|
||||
# architecture and scale from the checkpoint, so swapping the URL is enough to change it.
|
||||
_MODEL_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"
|
||||
_MODEL_FILENAME = "RealESRGAN_x2plus.pth"
|
||||
|
||||
_model: Any = None # lazy singleton (spandrel ImageModelDescriptor)
|
||||
_model_device: str = "cpu"
|
||||
_lock = threading.Lock()
|
||||
|
||||
|
||||
def is_available() -> bool:
|
||||
"""True if the ``esrgan`` extra (spandrel + torch) is importable."""
|
||||
from .optional_deps import module_available
|
||||
|
||||
return module_available("spandrel", "torch")
|
||||
|
||||
|
||||
def _model_cache_path() -> Path:
|
||||
"""Path the weights are cached at (the torch.hub checkpoints dir)."""
|
||||
import torch
|
||||
|
||||
cache_dir = Path(torch.hub.get_dir()) / "checkpoints"
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
return cache_dir / _MODEL_FILENAME
|
||||
|
||||
|
||||
def _get_model(device: str) -> Any:
|
||||
"""Load the Real-ESRGAN model once (downloading the weights on first use)."""
|
||||
global _model, _model_device
|
||||
if _model is not None and _model_device == device:
|
||||
return _model
|
||||
with _lock:
|
||||
if _model is None:
|
||||
import torch
|
||||
from spandrel import ImageModelDescriptor, ModelLoader
|
||||
|
||||
dst = _model_cache_path()
|
||||
if not dst.exists():
|
||||
logger.info("Downloading Real-ESRGAN weights to %s", dst)
|
||||
torch.hub.download_url_to_file(_MODEL_URL, str(dst), progress=False)
|
||||
model = ModelLoader().load_from_file(str(dst))
|
||||
if not isinstance(model, ImageModelDescriptor):
|
||||
raise RuntimeError(f"Unexpected spandrel model type: {type(model).__name__}")
|
||||
_model = model.eval()
|
||||
if _model_device != device:
|
||||
_model.to(device)
|
||||
_model_device = device
|
||||
return _model
|
||||
|
||||
|
||||
def scale() -> int:
|
||||
"""The model's native upscale factor (e.g. 2 for x2plus). Loads the model if needed."""
|
||||
return int(_get_model("cpu").scale)
|
||||
|
||||
|
||||
def upscale(image: NDArray[Any], device: str | None = None) -> NDArray[Any]:
|
||||
"""Upscale a BGR uint8 image by the model's native factor with Real-ESRGAN.
|
||||
|
||||
Returns a BGR uint8 array. Falls back to CPU if the requested device errors (an
|
||||
MPS/CUDA OOM or unsupported-op on the small pre-diffusion input), mirroring the
|
||||
diffusion engine's MPS->CPU fallback.
|
||||
|
||||
Raises:
|
||||
RuntimeError: if the ``esrgan`` extra is not installed (guard with
|
||||
``is_available()`` first).
|
||||
"""
|
||||
if not is_available():
|
||||
raise RuntimeError("Real-ESRGAN upscaler needs the 'esrgan' extra (spandrel). Install it or use Lanczos.")
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
target_device = (device or "cpu").lower()
|
||||
if target_device not in {"cpu", "mps", "cuda", "xpu"}:
|
||||
target_device = "cpu"
|
||||
rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
tensor = torch.from_numpy(rgb).permute(2, 0, 1).float().div(255.0).unsqueeze(0)
|
||||
|
||||
def _run(dev: str) -> NDArray[Any]:
|
||||
model = _get_model(dev)
|
||||
with torch.no_grad():
|
||||
out = model(tensor.to(dev))
|
||||
arr = out.clamp(0.0, 1.0).squeeze(0).permute(1, 2, 0).cpu().numpy() * 255.0
|
||||
return cv2.cvtColor(arr.round().astype(np.uint8), cv2.COLOR_RGB2BGR)
|
||||
|
||||
try:
|
||||
return _run(target_device)
|
||||
except Exception as e: # GPU OOM / unsupported op: fall back to CPU
|
||||
if target_device == "cpu":
|
||||
raise
|
||||
logger.warning("Real-ESRGAN on %s failed (%s); retrying on CPU", target_device, e)
|
||||
return _run("cpu")
|
||||
Reference in New Issue
Block a user