mirror of
https://github.com/D-Ogi/WatermarkRemover-AI.git
synced 2026-10-18 21:00:48 -05:00
Replace IOPaint with a verified standalone LaMA runtime (#88)
Replace the IOPaint dependency with a small verified TorchScript LaMA adapter, shared dependency manifests, reliable model preparation, and cross-platform application validation. Preserve crop/pad/composite behavior while handling empty masks and tiny inputs safely. Add integrity, offline, installer, image/video, and desktop backend coverage. Closes #48. Closes #49.
This commit is contained in:
D-Ogi authored and GitHub committed
1 parent
66d0c6b30b
commit
ae98ac0e10
24 files changed
+1417
-335
No files matched your search
@@ -36,7 +36,7 @@ jobs:
|
||||
- name: Check Python correctness
|
||||
run: python -m ruff check .
|
||||
- name: Compile application and tests without importing models
|
||||
run: python -m compileall -q remwm.py remwmgui.py utils.py tests
|
||||
run: python -m compileall -q remwm.py remwmgui.py utils.py lama_inpaint tests
|
||||
- name: Run CPU tests
|
||||
run: python -m pytest
|
||||
- name: Check Bash syntax
|
||||
@@ -57,3 +57,49 @@ jobs:
|
||||
$parseErrors | ForEach-Object { Write-Error $_ }
|
||||
exit 1
|
||||
}
|
||||
|
||||
application-tests:
|
||||
name: Application runtime (${{ matrix.os }})
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [ubuntu-latest, windows-latest, macos-latest]
|
||||
runs-on: ${{ matrix.os }}
|
||||
timeout-minutes: 25
|
||||
env:
|
||||
MPLBACKEND: Agg
|
||||
HF_HUB_OFFLINE: "1"
|
||||
TRANSFORMERS_OFFLINE: "1"
|
||||
steps:
|
||||
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
|
||||
with:
|
||||
persist-credentials: false
|
||||
- uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
- name: Install CPU PyTorch (Linux/Windows)
|
||||
if: runner.os != 'macOS'
|
||||
run: python -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
|
||||
- name: Install PyTorch (macOS)
|
||||
if: runner.os == 'macOS'
|
||||
run: python -m pip install torch torchvision
|
||||
- name: Install full application with the normal resolver
|
||||
run: python -m pip install -r requirements.txt pytest matplotlib
|
||||
- name: Verify full application environment
|
||||
run: |
|
||||
python -m pip check
|
||||
python -c "import remwm, webview, yaml, psutil; import importlib.util; assert importlib.util.find_spec('iopaint') is None"
|
||||
python remwm.py --help
|
||||
- name: Exercise application and TorchScript adapter without model downloads
|
||||
run: python -m pytest
|
||||
- name: Install Linux desktop system libraries
|
||||
if: runner.os == 'Linux'
|
||||
run: |
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y libegl1 libgl1 libxkbcommon-x11-0 libxcb-cursor0 libxcb-icccm4 libxcb-keysyms1 libxcb-shape0 libxcb-xinerama0 libxcb-randr0 libnss3 libasound2t64
|
||||
- name: Check installed desktop backend (Linux)
|
||||
if: runner.os == 'Linux'
|
||||
run: xvfb-run -a python scripts/check_webview.py
|
||||
- name: Check installed desktop backend (Windows/macOS)
|
||||
if: runner.os != 'Linux'
|
||||
run: python scripts/check_webview.py
|
||||
+4
-2
@@ -20,7 +20,9 @@ python -m compileall -q remwm.py remwmgui.py utils.py tests
|
||||
|
||||
CI runs these checks on Windows, Linux, and macOS. It also parses the PowerShell installer and checks Bash syntax and line endings. The lint baseline catches syntax errors and selected correctness errors; it does not require reformatting existing application code.
|
||||
|
||||
These checks use CPU fixtures. They do not install the full application, launch its GUI, download model weights, or prove CUDA/MPS compatibility. The test dependencies are separate from the application dependencies. A successful `pip check` in the test environment only validates that environment.
|
||||
These lightweight checks use CPU fixtures. They do not install the full application, launch its GUI, download model weights, or prove CUDA/MPS compatibility. The test dependencies are separate from the application dependencies. A successful `pip check` in the test environment only validates that environment.
|
||||
|
||||
A separate application CI job installs `requirements.txt` on all three platforms with the normal resolver, runs `pip check` in that environment, verifies application imports and CLI help, and tests the adapter using a generated TorchScript fixture. It also checks image and short-video routing without downloading detection models. This job does not prove a GUI interaction or real-model/GPU inference. See [LaMA runtime validation](docs/lama-runtime.md#validation).
|
||||
|
||||
## Evidence required by change type
|
||||
|
||||
@@ -45,7 +47,7 @@ GPU/model integration tests require a separate trusted environment and must not
|
||||
|
||||
Automated review provides suggestions; maintainers verify findings against the code and test evidence. A bot's approval does not replace application testing. Resolve applicable review findings or explain why they do not apply. Maintainers decide when to merge.
|
||||
|
||||
Once the new CI jobs have passed reliably, maintainers can require their status checks through repository rules. Adding this workflow alone does not enable branch protection. Full installation and GUI automation remain follow-up work while the application's dependency and backend setup issues are being resolved.
|
||||
Once the new CI jobs have passed reliably, maintainers can require their status checks through repository rules. Adding this workflow alone does not enable branch protection. Full-manifest installation checks run in the application job; installer-script execution and GUI interaction still require separate validation on affected platforms.
|
||||
|
||||
## Dependency updates
|
||||
|
||||
|
||||
@@ -162,3 +162,10 @@ This project is licensed under the MIT License. See the [LICENSE](LICENSE) file
|
||||
[](https://www.star-history.com/#D-Ogi/WatermarkRemover-AI&type=date&legend=top-left)
|
||||
|
||||
|
||||
|
||||
## LaMA runtime and installation troubleshooting
|
||||
|
||||
LaMA inpainting runs directly from verified TorchScript weights; IOPaint is no
|
||||
longer required. See [runtime, cache and migration instructions](docs/lama-runtime.md)
|
||||
when upgrading an older installation. Use a fresh application environment and check
|
||||
it with `python -m pip check` to avoid retaining old IOPaint dependency conflicts.
|
||||
@@ -0,0 +1,28 @@
|
||||
# Third-party notices
|
||||
|
||||
## LaMA processing adapter
|
||||
|
||||
The crop, symmetric padding, tensor normalization and compositing behavior in
|
||||
`lama_inpaint` is adapted from [Sanster/IOPaint](https://github.com/Sanster/IOPaint),
|
||||
revision `61a759fb3f332bacdce8b2813f4837495c9b86e0`, specifically
|
||||
`iopaint/model/lama.py`, `iopaint/model/base.py` and `iopaint/helper.py`.
|
||||
The upstream Apache-2.0 license is reproduced in
|
||||
[licenses/IOPaint-Apache-2.0.txt](licenses/IOPaint-Apache-2.0.txt).
|
||||
|
||||
Changes: standalone RGB prediction/BGR output adapter; input validation;
|
||||
non-mutating compositing; empty-mask bypass; soft-mask crop selection;
|
||||
verified atomic model downloads; removal of the IOPaint registry, schemas,
|
||||
CLI, diffusion models and their runtime dependencies.
|
||||
|
||||
## Model weights
|
||||
|
||||
Weights are downloaded separately from the existing upstream
|
||||
[big-lama.pt release](https://github.com/Sanster/models/releases/tag/add_big_lama).
|
||||
They are not included in this repository. The artifact is 205669692 bytes;
|
||||
SHA-256: `344c77bbcb158f17dd143070d1e789f38a66c04202311ae3a258ef66667a9ea9`.
|
||||
This digest was calculated from the artifact matching IOPaint's pinned MD5
|
||||
`e3aa4aaa15225a33ec84f9f4bc47e500`; the legacy release has no GitHub SHA-256
|
||||
asset digest. The runtime uses the pinned SHA-256, not MD5 or an environment
|
||||
variable supplied checksum. This is an integrity pin, not an upstream signature.
|
||||
For model provenance and terms, see [LaMA](https://github.com/advimman/lama)
|
||||
and the artifact publisher. Code licensing does not replace model/data terms.
|
||||
@@ -0,0 +1,134 @@
|
||||
# Standalone LaMA runtime
|
||||
|
||||
The application loads the existing `big-lama.pt` TorchScript model directly.
|
||||
IOPaint, its model registry/CLI, diffusion backends, and request schemas are no
|
||||
longer needed. Florence-2 detection, desktop themes and translations retain their
|
||||
existing behavior.
|
||||
|
||||
## Installation and migration
|
||||
|
||||
Use a fresh application environment when upgrading an installation that contains
|
||||
IOPaint. Leaving that package installed can retain conflicting requirements even
|
||||
though the application no longer imports it. The installers do not remove packages
|
||||
from another environment and do not suppress dependency failures.
|
||||
|
||||
For a manual installation with Python 3.12:
|
||||
|
||||
```sh
|
||||
python -m venv .venv-app
|
||||
# Linux/macOS: . .venv-app/bin/activate
|
||||
# Windows PowerShell: .venv-app\Scripts\Activate.ps1
|
||||
python -m pip install -r requirements.txt
|
||||
python -m pip check
|
||||
python -m lama_inpaint download
|
||||
python remwm.py --help
|
||||
```
|
||||
|
||||
For a headless CLI installation, use `requirements-core.txt` instead. Select a
|
||||
matching torch/torchvision build for your device before installing requirements;
|
||||
the generic manifest does not guarantee a CUDA build. Linux desktop installation
|
||||
uses pywebview's Qt extra, Windows keeps its native backend, and macOS keeps its
|
||||
native backend. Linux still needs a graphical session and the Qt system libraries. On Ubuntu
|
||||
24.04, prepare the desktop libraries with:
|
||||
|
||||
```sh
|
||||
sudo apt-get install libegl1 libgl1 libxkbcommon-x11-0 libxcb-cursor0 libxcb-icccm4 libxcb-keysyms1 libxcb-shape0 libxcb-xinerama0 libxcb-randr0 libnss3 libasound2t64
|
||||
```
|
||||
|
||||
Package names may differ on other distributions. The installer does not silently
|
||||
install system packages or elevate privileges.
|
||||
|
||||
The standalone runtime removes the IOPaint `imghdr` import path. This does not by
|
||||
itself certify every dependency or installer on Python 3.13+. The application CI
|
||||
currently validates Python 3.12.
|
||||
|
||||
## Model integrity and offline use
|
||||
|
||||
All three installers and runtime loading use the same model store. It reuses
|
||||
`$TORCH_HOME/hub/checkpoints/big-lama.pt`, or otherwise
|
||||
`$XDG_CACHE_HOME/torch/hub/checkpoints/big-lama.pt` (default `~/.cache/torch/...`).
|
||||
An existing valid IOPaint cache is reused without downloading it again. A custom
|
||||
`torch.hub.set_dir()` made by unrelated Python code is not consulted.
|
||||
|
||||
```sh
|
||||
python -m lama_inpaint download
|
||||
python -m lama_inpaint download --offline
|
||||
python -m lama_inpaint download --cache-dir /path/to/checkpoints
|
||||
```
|
||||
|
||||
`--cache-dir` applies to that preparation command; set `TORCH_HOME` for the normal
|
||||
application to use a nondefault cache. Offline verification fails if the file is
|
||||
missing or corrupt and never opens a network connection. Programmatic callers
|
||||
can use `LamaInpaint(device="cpu", download=False)` for strict offline loading.
|
||||
|
||||
The artifact URL, exact byte count and SHA-256 are pinned in
|
||||
`lama_inpaint/model_store.py`; see [provenance](../THIRD_PARTY_NOTICES.md).
|
||||
Cached weights are verified before deserialization. Downloads have read/time/size
|
||||
limits, a cross-process lock, a unique temporary file and atomic replacement after
|
||||
verification. Interrupted or invalid downloads never become the cached model and
|
||||
do not delete an existing artifact. An invalid cache is replaced only after a new
|
||||
download passes verification. No user-provided URL or checksum environment variable
|
||||
can bypass the pinned artifact check.
|
||||
|
||||
## Restricted networks
|
||||
|
||||
The China option changes the package and Florence-2 mirrors. LaMA still uses the
|
||||
pinned GitHub artifact; there is no unverified model mirror or URL override.
|
||||
If GitHub is unavailable, preseed the cache before running setup:
|
||||
|
||||
1. On a machine with GitHub access and these application packages installed, run
|
||||
`python -m lama_inpaint download --cache-dir model-transfer`.
|
||||
2. Transfer `model-transfer/big-lama.pt` to the target machine. Place it at
|
||||
`TORCH_HOME/hub/checkpoints/big-lama.pt` if you set `TORCH_HOME`, otherwise at
|
||||
`~/.cache/torch/hub/checkpoints/big-lama.pt`. If `XDG_CACHE_HOME` is set, use
|
||||
`XDG_CACHE_HOME/torch/hub/checkpoints/big-lama.pt` instead of `~/.cache/...`.
|
||||
Keep the cache environment variables consistent when running setup and the app.
|
||||
3. Run setup normally with the chosen package mirror. Every installer checks the
|
||||
local artifact before attempting any model network request. Valid preseeded
|
||||
weights are reused; missing or incorrect weights still cause setup to fail if
|
||||
GitHub cannot be reached. Do not bypass the checksum check.
|
||||
4. In the installed application environment, confirm the cache without network
|
||||
access using `python -m lama_inpaint download --offline`.
|
||||
|
||||
Linux installs CPU and CUDA PyTorch builds from their respective official wheel
|
||||
indexes, including when a different mirror is selected for other packages. The
|
||||
CUDA branch currently selects CUDA 12.4; this is not a promise of support for every
|
||||
GPU/driver generation. Hardware requiring another build still needs an explicitly
|
||||
selected compatible torch/torchvision pair.
|
||||
|
||||
## Processing contract
|
||||
|
||||
- Input: nonempty `uint8` RGB array `(H, W, 3)` and matching `uint8` mask `(H, W)`.
|
||||
- Output: independent `uint8` BGR array, preserving the application's existing API.
|
||||
- Values above zero select model pixels; mask values divided by 255 determine the
|
||||
final composite. Zero-mask pixels are preserved exactly.
|
||||
- Images above 800 pixels on the longest side use connected-region crops with the
|
||||
existing 64-pixel context margin. Crops are predicted from the original image,
|
||||
including overlapping context regions.
|
||||
- Symmetric padding aligns dimensions to multiples of 8. Tiny inputs are padded
|
||||
to at least 32 pixels per dimension, then cropped back; they are never resized.
|
||||
- Empty masks bypass inference. Caller-owned image/mask arrays are never modified.
|
||||
- CPU and CUDA are supported. An unavailable requested CUDA device fails explicitly;
|
||||
MPS is rejected until its operators and results are validated.
|
||||
|
||||
The old `resize_limit=1600` belonged to an unused RESIZE strategy; the application
|
||||
selected CROP. It was not a hard cap on a full-image mask. The new adapter likewise
|
||||
does not resize large masks silently. Memory requirements still depend on crop size.
|
||||
|
||||
Intentional differences from the previous wrapper: no input mutation, no model
|
||||
call for an empty mask, safe padding for tiny images, and inclusion of soft masks
|
||||
below 128 when choosing large-image crops. Standard binary-mask processing retains
|
||||
the previous crop/pad/composite behavior.
|
||||
|
||||
## Validation
|
||||
|
||||
The lightweight suite exercises processing and download integrity without PyTorch
|
||||
or real model weights. The separate application job installs the complete manifest
|
||||
on Windows/Linux/macOS, runs `pip check`, checks imports/CLI help, and executes a
|
||||
small generated TorchScript model plus image/video routing tests. Neither job
|
||||
downloads production weights or certifies GPU execution.
|
||||
|
||||
Changes to model execution also need a separately recorded real-weight comparison:
|
||||
model digest, runtime versions, device, fixture dimensions/masks and differences
|
||||
from the reference. Keep fixtures synthetic or redistributable. GPU coverage must
|
||||
name the tested device; it does not imply universal GPU or MPS compatibility.
|
||||
@@ -0,0 +1,5 @@
|
||||
"""Small LaMA inference adapter; no IOPaint runtime dependency."""
|
||||
|
||||
from .processing import InpaintProcessor
|
||||
|
||||
__all__ = ["InpaintProcessor"]
|
||||
@@ -0,0 +1,30 @@
|
||||
"""Explicit model preparation: python -m lama_inpaint download [--offline]."""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
from .model_store import ModelError, ensure_model
|
||||
|
||||
|
||||
def main():
|
||||
"""Verify/download the fixed artifact, returning a nonzero status on failure."""
|
||||
parser = argparse.ArgumentParser(description="Prepare verified LaMA model weights")
|
||||
parser.add_argument("command", choices=["download"])
|
||||
parser.add_argument("--cache-dir")
|
||||
parser.add_argument(
|
||||
"--offline",
|
||||
action="store_true",
|
||||
help="Verify cached weights without network access",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
try:
|
||||
path = ensure_model(args.cache_dir, download=not args.offline)
|
||||
except (ModelError, OSError) as exc:
|
||||
print(f"LaMA model preparation failed: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
print(f"LaMA model verified: {path}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,114 @@
|
||||
"""Verified, atomic download and reuse of the fixed LaMA TorchScript artifact."""
|
||||
|
||||
import hashlib
|
||||
from http.client import HTTPException
|
||||
import os
|
||||
from pathlib import Path
|
||||
import tempfile
|
||||
import time
|
||||
from urllib.request import urlopen
|
||||
|
||||
from filelock import FileLock, Timeout
|
||||
|
||||
MODEL_URL = (
|
||||
"https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt"
|
||||
)
|
||||
MODEL_SHA256 = "344c77bbcb158f17dd143070d1e789f38a66c04202311ae3a258ef66667a9ea9"
|
||||
MODEL_SIZE = 205669692
|
||||
MODEL_NAME = "big-lama.pt"
|
||||
|
||||
|
||||
class ModelError(RuntimeError):
|
||||
"""A model could not be obtained or verified; no unverified weights are loaded."""
|
||||
|
||||
|
||||
def default_cache_dir():
|
||||
"""Honor the standard Torch cache variables and reuse existing LaMA weights."""
|
||||
torch_home = os.environ.get("TORCH_HOME")
|
||||
if torch_home:
|
||||
return Path(torch_home).expanduser() / "hub" / "checkpoints"
|
||||
cache = Path(os.environ.get("XDG_CACHE_HOME", "~/.cache")).expanduser()
|
||||
return cache / "torch" / "hub" / "checkpoints"
|
||||
|
||||
|
||||
def verify_stream(stream):
|
||||
"""Check an open artifact's exact size and SHA-256, then rewind it."""
|
||||
stream.seek(0)
|
||||
digest = hashlib.sha256()
|
||||
size = 0
|
||||
while chunk := stream.read(1024 * 1024):
|
||||
size += len(chunk)
|
||||
if size > MODEL_SIZE:
|
||||
raise ModelError("LaMA model exceeds the expected artifact size.")
|
||||
digest.update(chunk)
|
||||
stream.seek(0)
|
||||
if size != MODEL_SIZE or digest.hexdigest() != MODEL_SHA256:
|
||||
raise ModelError("LaMA model checksum or size mismatch; refusing to load it.")
|
||||
|
||||
|
||||
def _is_valid(path):
|
||||
"""Return whether an existing artifact matches the pinned size and digest."""
|
||||
try:
|
||||
with path.open("rb") as stream:
|
||||
verify_stream(stream)
|
||||
return True
|
||||
except (FileNotFoundError, ModelError):
|
||||
return False
|
||||
|
||||
|
||||
def ensure_model(cache_dir=None, *, download=True):
|
||||
"""Return verified weights, downloading atomically when allowed.
|
||||
|
||||
Downloads use a unique temporary file and a per-artifact cross-process lock.
|
||||
A failed transfer never replaces an existing file. Offline mode never opens
|
||||
a network connection. The URL and expected digest are deliberately fixed.
|
||||
"""
|
||||
cache = Path(cache_dir) if cache_dir is not None else default_cache_dir()
|
||||
target = cache / MODEL_NAME
|
||||
if not download:
|
||||
if not _is_valid(target):
|
||||
raise ModelError(
|
||||
f"No verified LaMA model at {target}. Run python -m lama_inpaint download first."
|
||||
)
|
||||
return target
|
||||
cache.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
with FileLock(str(target) + ".lock", timeout=600):
|
||||
if _is_valid(target):
|
||||
return target
|
||||
_download(target)
|
||||
except Timeout as exc:
|
||||
raise ModelError("Timed out waiting for another LaMA model download.") from exc
|
||||
return target
|
||||
|
||||
|
||||
def _download(target):
|
||||
"""Publish only a complete verified transfer; retain the old file on failure."""
|
||||
temporary = None
|
||||
try:
|
||||
with tempfile.NamedTemporaryFile(
|
||||
dir=target.parent, prefix=".lama-", suffix=".part", delete=False
|
||||
) as output:
|
||||
temporary = Path(output.name)
|
||||
deadline = time.monotonic() + 600
|
||||
size = 0
|
||||
with urlopen(MODEL_URL, timeout=30) as response:
|
||||
while chunk := response.read(1024 * 1024):
|
||||
size += len(chunk)
|
||||
if size > MODEL_SIZE or time.monotonic() > deadline:
|
||||
raise ModelError(
|
||||
"LaMA download exceeded its size or time limit."
|
||||
)
|
||||
output.write(chunk)
|
||||
output.flush()
|
||||
os.fsync(output.fileno())
|
||||
with temporary.open("rb") as stream:
|
||||
verify_stream(stream)
|
||||
os.replace(temporary, target)
|
||||
except (OSError, ValueError, HTTPException) as exc:
|
||||
raise ModelError(
|
||||
f"Could not download LaMA weights: {exc}. Retry python -m lama_inpaint download."
|
||||
) from exc
|
||||
finally:
|
||||
if temporary is not None:
|
||||
temporary.unlink(missing_ok=True)
|
||||
@@ -0,0 +1,103 @@
|
||||
"""LaMA crop/pad/composite processing with an injectable RGB predictor.
|
||||
|
||||
Adapted from IOPaint's LaMa/InpaintModel and image helpers (Apache-2.0).
|
||||
See THIRD_PARTY_NOTICES.md for the source revision and license.
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
class InpaintProcessor:
|
||||
"""Process uint8 RGB images and uint8 masks; return independent uint8 BGR.
|
||||
|
||||
The predictor receives padded RGB and single-channel masks and returns RGB.
|
||||
A positive mask value selects model input; its 0..255 value controls final
|
||||
blending. Large images use the existing 800-pixel crop trigger and 64-pixel
|
||||
context margin. Predictions always see the original input, including when
|
||||
context crops overlap. Neither caller-owned input is modified.
|
||||
"""
|
||||
|
||||
def __init__(self, predictor, *, crop_trigger=800, crop_margin=64):
|
||||
"""Bind an RGB predictor and validate the positive crop trigger and margin."""
|
||||
if not isinstance(crop_trigger, int) or crop_trigger < 1:
|
||||
raise ValueError("crop_trigger must be a positive integer")
|
||||
if not isinstance(crop_margin, int) or crop_margin < 1:
|
||||
raise ValueError("crop_margin must be a positive integer")
|
||||
self.predictor = predictor
|
||||
self.crop_trigger = crop_trigger
|
||||
self.crop_margin = crop_margin
|
||||
|
||||
def __call__(self, image, mask):
|
||||
"""Validate RGB/mask arrays and return a new BGR composite without input mutation."""
|
||||
if (
|
||||
not isinstance(image, np.ndarray)
|
||||
or image.dtype != np.uint8
|
||||
or image.ndim != 3
|
||||
or image.shape[2] != 3
|
||||
):
|
||||
raise ValueError("image must be an HxWx3 uint8 RGB array")
|
||||
if (
|
||||
not isinstance(mask, np.ndarray)
|
||||
or mask.dtype != np.uint8
|
||||
or mask.shape != image.shape[:2]
|
||||
):
|
||||
raise ValueError("mask must be an HxW uint8 array matching the image")
|
||||
height, width = mask.shape
|
||||
if not height or not width:
|
||||
raise ValueError("image and mask must not be empty")
|
||||
result = image[:, :, ::-1].copy()
|
||||
if not mask.any():
|
||||
return result
|
||||
if max(height, width) <= self.crop_trigger:
|
||||
return self._forward(image, mask)
|
||||
binary = (mask > 0).astype(np.uint8)
|
||||
contours, _ = cv2.findContours(
|
||||
binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
|
||||
)
|
||||
for contour in contours:
|
||||
x, y, w, h = cv2.boundingRect(contour)
|
||||
left, right = self._bounds(x, w, width)
|
||||
top, bottom = self._bounds(y, h, height)
|
||||
crop = self._forward(
|
||||
image[top:bottom, left:right], mask[top:bottom, left:right]
|
||||
)
|
||||
result[top:bottom, left:right] = crop
|
||||
return result
|
||||
|
||||
def _bounds(self, start, length, limit):
|
||||
# Preserve IOPaint's centered crop sizing, including odd-sized boxes.
|
||||
"""Return a clipped context interval, preserving legacy odd-box centering."""
|
||||
center = (2 * start + length) // 2
|
||||
half = (length + 2 * self.crop_margin) // 2
|
||||
lower, upper = center - half, center + half
|
||||
left, right = max(lower, 0), min(upper, limit)
|
||||
if lower < 0:
|
||||
right += -lower
|
||||
if upper > limit:
|
||||
left -= upper - limit
|
||||
return max(left, 0), min(right, limit)
|
||||
|
||||
def _forward(self, image, mask):
|
||||
"""Symmetrically pad one crop, predict RGB, and composite its original-size BGR result."""
|
||||
height, width = mask.shape
|
||||
# Reflection padding inside the network needs more than a single feature
|
||||
# pixel after downsampling. Keep tiny inputs safe without resizing them.
|
||||
padded_height = max(32, ((height + 7) // 8) * 8)
|
||||
padded_width = max(32, ((width + 7) // 8) * 8)
|
||||
padding = ((0, padded_height - height), (0, padded_width - width))
|
||||
padded_image = np.pad(image, (*padding, (0, 0)), mode="symmetric")
|
||||
padded_mask = np.pad(mask, padding, mode="symmetric")
|
||||
prediction = self.predictor(padded_image, padded_mask)
|
||||
if (
|
||||
not isinstance(prediction, np.ndarray)
|
||||
or prediction.shape != padded_image.shape
|
||||
or prediction.dtype != np.uint8
|
||||
):
|
||||
raise ValueError(
|
||||
"LaMA predictor must return a uint8 RGB array matching its padded input"
|
||||
)
|
||||
prediction = prediction[:height, :width, ::-1]
|
||||
# float64 matches the previous IOPaint composite and truncation behavior.
|
||||
alpha = mask[:, :, None] / 255.0
|
||||
return (prediction * alpha + image[:, :, ::-1] * (1 - alpha)).astype(np.uint8)
|
||||
@@ -0,0 +1,48 @@
|
||||
"""TorchScript LaMA execution, independent of detection and GUI frameworks."""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .model_store import ModelError, ensure_model, verify_stream
|
||||
from .processing import InpaintProcessor
|
||||
|
||||
|
||||
class LamaInpaint(InpaintProcessor):
|
||||
"""Load verified LaMA on CPU/CUDA and preserve the application's BGR API."""
|
||||
|
||||
def __init__(self, device="cpu", *, cache_dir=None, download=True):
|
||||
"""Verify the pinned artifact and load an evaluation-only CPU or CUDA model."""
|
||||
self.device = torch.device(device)
|
||||
if self.device.type not in {"cpu", "cuda"}:
|
||||
raise ValueError(
|
||||
"LaMA currently supports CPU and CUDA; MPS is not validated"
|
||||
)
|
||||
if self.device.type == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError("CUDA was requested for LaMA but PyTorch cannot use it")
|
||||
path = ensure_model(cache_dir, download=download)
|
||||
try:
|
||||
# Verify and load through the same handle: never deserialize a replaced path.
|
||||
with path.open("rb") as stream:
|
||||
verify_stream(stream)
|
||||
self.model = (
|
||||
torch.jit.load(stream, map_location="cpu").eval().to(self.device)
|
||||
)
|
||||
except (OSError, RuntimeError) as exc:
|
||||
raise ModelError(
|
||||
f"Could not load verified LaMA weights on {self.device}: {exc}"
|
||||
) from exc
|
||||
super().__init__(self._predict)
|
||||
|
||||
@torch.inference_mode()
|
||||
def _predict(self, image, mask):
|
||||
"""Run normalized NCHW tensors and validate finite RGB output before quantization."""
|
||||
rgb = np.ascontiguousarray(image.transpose(2, 0, 1), dtype=np.float32) / 255.0
|
||||
selected = np.ascontiguousarray(mask[None] > 0, dtype=np.int64)
|
||||
result = self.model(
|
||||
torch.from_numpy(rgb)[None].to(self.device),
|
||||
torch.from_numpy(selected)[None].to(self.device),
|
||||
)
|
||||
if result.shape != (1, 3, *mask.shape) or not torch.isfinite(result).all():
|
||||
raise RuntimeError("LaMA returned invalid dimensions or non-finite pixels")
|
||||
rgb = result[0].permute(1, 2, 0).cpu().numpy()
|
||||
return np.clip(rgb * 255, 0, 255).astype(np.uint8)
|
||||
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -5,14 +5,8 @@ import cv2
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
# Monkey-patch: cached_download was removed in huggingface_hub 0.24, add compatibility shim
|
||||
import huggingface_hub
|
||||
if not hasattr(huggingface_hub, 'cached_download'):
|
||||
huggingface_hub.cached_download = huggingface_hub.hf_hub_download
|
||||
|
||||
from transformers import AutoProcessor, Florence2ForConditionalGeneration
|
||||
from iopaint.model_manager import ModelManager
|
||||
from iopaint.schema import HDStrategy, LDMSampler, InpaintRequest as Config
|
||||
from lama_inpaint.runtime import LamaInpaint
|
||||
import torch
|
||||
from torch.nn import Module
|
||||
import tqdm
|
||||
@@ -29,43 +23,11 @@ except ImportError:
|
||||
MatLike = np.ndarray
|
||||
|
||||
|
||||
def download_lama_model():
|
||||
"""Download LaMA model using iopaint."""
|
||||
logger.info("Downloading LaMA model... (this may take a few minutes)")
|
||||
print("Downloading LaMA model (~196MB)... Please wait.")
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-m", "iopaint", "download", "--model", "lama"],
|
||||
capture_output=False, # Show download progress
|
||||
text=True
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
logger.error("Failed to download LaMA model")
|
||||
return False
|
||||
|
||||
logger.info("LaMA model downloaded successfully")
|
||||
print("LaMA model downloaded!")
|
||||
return True
|
||||
|
||||
|
||||
def load_lama_model(device):
|
||||
"""Load LaMA model, downloading if necessary."""
|
||||
try:
|
||||
return ModelManager(name="lama", device=device)
|
||||
except NotImplementedError as e:
|
||||
if "Unsupported model: lama" in str(e):
|
||||
print("LaMA model not available, attempting to download...")
|
||||
if download_lama_model():
|
||||
# Re-import to refresh model registry
|
||||
import importlib
|
||||
import iopaint.model
|
||||
importlib.reload(iopaint.model)
|
||||
# Try again
|
||||
return ModelManager(name="lama", device=device)
|
||||
else:
|
||||
raise RuntimeError("Failed to download LaMA model. Please run manually: python\\python.exe -m iopaint download --model lama")
|
||||
raise
|
||||
"""Load checksum-verified LaMA weights using the standalone adapter."""
|
||||
logger.info("Loading LaMA (verifying cached weights or downloading if missing)")
|
||||
return LamaInpaint(device)
|
||||
|
||||
|
||||
class TaskType(str, Enum):
|
||||
OPEN_VOCAB_DETECTION = "<OPEN_VOCABULARY_DETECTION>"
|
||||
@@ -156,16 +118,9 @@ def detect_only(image: MatLike, model: Florence2ForConditionalGeneration, proces
|
||||
|
||||
return results
|
||||
|
||||
def process_image_with_lama(image: MatLike, mask: MatLike, model_manager: ModelManager):
|
||||
config = Config(
|
||||
ldm_steps=50,
|
||||
ldm_sampler=LDMSampler.ddim,
|
||||
hd_strategy=HDStrategy.CROP,
|
||||
hd_strategy_crop_margin=64,
|
||||
hd_strategy_crop_trigger_size=800,
|
||||
hd_strategy_resize_limit=1600,
|
||||
)
|
||||
result = model_manager(image, mask, config)
|
||||
def process_image_with_lama(image: MatLike, mask: MatLike, model_manager: LamaInpaint):
|
||||
"""Return the shared LaMA adapter result as a uint8 BGR image."""
|
||||
result = model_manager(image, mask)
|
||||
|
||||
if result.dtype in [np.float64, np.float32]:
|
||||
result = np.clip(result, 0, 255).astype(np.uint8)
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
# CLI and model runtime. Install a matching torch/torchvision build for your device first.
|
||||
torch>=2.4.0
|
||||
torchvision>=0.19.0
|
||||
# Native Florence-2 API; pre-4.56 releases do not provide the required model class.
|
||||
transformers>=4.56.0,<5
|
||||
numpy>=1.26,<2
|
||||
opencv-python-headless>=4.8.0,<4.12.0
|
||||
Pillow>=10.0.0
|
||||
filelock>=3.16
|
||||
loguru
|
||||
click
|
||||
tqdm
|
||||
@@ -4,3 +4,5 @@ ruff==0.16.7
|
||||
numpy==2.4.6
|
||||
Pillow==12.3.0
|
||||
matplotlib==3.10.8
|
||||
opencv-python-headless==4.11.0.86
|
||||
filelock==3.32.3
|
||||
+4
-24
@@ -1,27 +1,7 @@
|
||||
# Core AI/ML - CUDA 12.4 versions
|
||||
--extra-index-url https://download.pytorch.org/whl/cu124
|
||||
torch>=2.4.0
|
||||
torchvision>=0.19.0
|
||||
transformers>=4.50.0
|
||||
diffusers>=0.30.0
|
||||
numpy<2
|
||||
-r requirements-core.txt
|
||||
|
||||
# Image/Video processing
|
||||
opencv-python-headless>=4.8.0,<4.12.0
|
||||
Pillow>=10.0.0
|
||||
|
||||
# Inpainting core (iopaint installed separately with --no-deps to avoid resolver conflicts)
|
||||
|
||||
# GUI (PyWebview for brainrot HTML UI)
|
||||
pywebview>=4.0
|
||||
|
||||
# Utilities
|
||||
loguru
|
||||
click
|
||||
tqdm
|
||||
# Desktop UI. Linux needs a concrete backend; Qt wheels avoid a GTK source build.
|
||||
pywebview>=4.0; sys_platform != "linux"
|
||||
pywebview[qt]>=4.0; sys_platform == "linux"
|
||||
psutil
|
||||
pyyaml
|
||||
|
||||
# qt and gtk gui
|
||||
PyGObject
|
||||
qtpy
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Check that the installed desktop backend can call Python from JavaScript."""
|
||||
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
|
||||
import webview
|
||||
|
||||
|
||||
class Bridge:
|
||||
"""Minimal API used only by this installation check."""
|
||||
|
||||
def ping(self):
|
||||
"""Return a value for the browser to expose after resolving its API call."""
|
||||
return "bridge-ok"
|
||||
|
||||
|
||||
completed = threading.Event()
|
||||
finished = threading.Event()
|
||||
window = webview.create_window(
|
||||
"Backend verification",
|
||||
html="""<html><body><script>
|
||||
window.addEventListener('pywebviewready', async () => {
|
||||
window.bridgeResult = await window.pywebview.api.ping();
|
||||
});
|
||||
</script></body></html>""",
|
||||
js_api=Bridge(),
|
||||
hidden=True,
|
||||
)
|
||||
|
||||
|
||||
def check_bridge():
|
||||
"""Close only after the browser has received its Python API response."""
|
||||
deadline = time.monotonic() + 45
|
||||
while time.monotonic() < deadline:
|
||||
if window.evaluate_js("window.bridgeResult || null") == "bridge-ok":
|
||||
completed.set()
|
||||
break
|
||||
time.sleep(0.1)
|
||||
window.destroy()
|
||||
|
||||
|
||||
def watchdog():
|
||||
"""Fail even if startup, a native JS call or window teardown hangs."""
|
||||
if not finished.wait(60):
|
||||
print("Backend bridge/startup/teardown timed out", flush=True)
|
||||
os._exit(1)
|
||||
|
||||
|
||||
threading.Thread(target=watchdog, daemon=True).start()
|
||||
webview.start(check_bridge)
|
||||
finished.set()
|
||||
if not completed.is_set():
|
||||
raise SystemExit("Backend closed before completing its bridge check")
|
||||
print("Desktop backend JavaScript/Python bridge passed")
|
||||
@@ -82,6 +82,13 @@ if not exist "%PYTHON_EXE%" (
|
||||
echo [OK] Python found
|
||||
)
|
||||
|
||||
:: Include the application root for embedded Python package imports.
|
||||
set PTH_FILE=%PYTHON_DIR%\python312._pth
|
||||
if exist "%PTH_FILE%" (
|
||||
findstr /l /x /c:".." "%PTH_FILE%" >nul
|
||||
if errorlevel 1 echo ..>> "%PTH_FILE%"
|
||||
)
|
||||
|
||||
echo.
|
||||
echo [*] Installing dependencies...
|
||||
echo This takes 5-10 minutes. Please wait...
|
||||
@@ -97,9 +104,9 @@ if "%CHINA_MODE%"=="1" (
|
||||
:: Install base dependencies
|
||||
echo [*] Installing base packages...
|
||||
if "%CHINA_MODE%"=="1" (
|
||||
"%PYTHON_EXE%" -m pip install --upgrade -r requirements.txt --no-cache-dir --use-deprecated=legacy-resolver %PIP_MIRROR%
|
||||
"%PYTHON_EXE%" -m pip install --upgrade -r requirements.txt --no-cache-dir %PIP_MIRROR%
|
||||
) else (
|
||||
"%PYTHON_EXE%" -m pip install --upgrade -r requirements.txt --no-cache-dir --use-deprecated=legacy-resolver
|
||||
"%PYTHON_EXE%" -m pip install --upgrade -r requirements.txt --no-cache-dir
|
||||
)
|
||||
if errorlevel 1 (
|
||||
echo [X] Failed to install base dependencies
|
||||
@@ -107,89 +114,28 @@ if errorlevel 1 (
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
:: Verify key packages
|
||||
"%PYTHON_EXE%" -c "import torch; import transformers; import webview; import cv2; print('OK')" >nul 2>&1
|
||||
:: Verify the full application environment, not only selected imports.
|
||||
"%PYTHON_EXE%" -m pip check
|
||||
if errorlevel 1 (
|
||||
echo [X] Failed to verify base packages
|
||||
echo [X] Dependency verification failed. Use a fresh application environment.
|
||||
pause
|
||||
exit /b 1
|
||||
)
|
||||
echo [OK] Base packages installed
|
||||
|
||||
:: Install iopaint without dependencies
|
||||
echo [*] Installing iopaint...
|
||||
if "%CHINA_MODE%"=="1" (
|
||||
"%PYTHON_EXE%" -m pip install --upgrade iopaint --no-deps --no-cache-dir %PIP_MIRROR%
|
||||
) else (
|
||||
"%PYTHON_EXE%" -m pip install --upgrade iopaint --no-deps --no-cache-dir
|
||||
)
|
||||
"%PYTHON_EXE%" -c "import remwm; import webview; import yaml; import psutil"
|
||||
if errorlevel 1 (
|
||||
echo [X] Failed to install iopaint
|
||||
pause
|
||||
exit /b 1
|
||||
)
|
||||
echo [OK] iopaint installed
|
||||
|
||||
:: Install iopaint dependencies manually
|
||||
echo [*] Installing iopaint dependencies...
|
||||
if "%CHINA_MODE%"=="1" (
|
||||
"%PYTHON_EXE%" -m pip install pydantic typer einops omegaconf easydict yacs --no-cache-dir %PIP_MIRROR%
|
||||
) else (
|
||||
"%PYTHON_EXE%" -m pip install pydantic typer einops omegaconf easydict yacs --no-cache-dir
|
||||
)
|
||||
if errorlevel 1 (
|
||||
echo [X] Failed to install iopaint dependencies
|
||||
echo [X] Failed to import application packages
|
||||
pause
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
:: Verify iopaint dependencies
|
||||
"%PYTHON_EXE%" -c "import pydantic; import typer; import einops; import omegaconf; import easydict; import yacs; print('OK')" >nul 2>&1
|
||||
:: Shared checksum-verified, atomic model download.
|
||||
echo [*] Preparing LaMA model (196MB)...
|
||||
if "%CHINA_MODE%"=="1" echo If GitHub is blocked, preseed the verified cache: docs/lama-runtime.md#restricted-networks
|
||||
"%PYTHON_EXE%" -m lama_inpaint download
|
||||
if errorlevel 1 (
|
||||
echo [!] iopaint dependencies verification failed, attempting reinstall...
|
||||
if "%CHINA_MODE%"=="1" (
|
||||
"%PYTHON_EXE%" -m pip install pydantic --no-cache-dir --force-reinstall %PIP_MIRROR%
|
||||
"%PYTHON_EXE%" -m pip install typer --no-cache-dir --force-reinstall %PIP_MIRROR%
|
||||
"%PYTHON_EXE%" -m pip install einops --no-cache-dir --force-reinstall %PIP_MIRROR%
|
||||
"%PYTHON_EXE%" -m pip install omegaconf --no-cache-dir --force-reinstall %PIP_MIRROR%
|
||||
"%PYTHON_EXE%" -m pip install easydict --no-cache-dir --force-reinstall %PIP_MIRROR%
|
||||
"%PYTHON_EXE%" -m pip install yacs --no-cache-dir --force-reinstall %PIP_MIRROR%
|
||||
) else (
|
||||
"%PYTHON_EXE%" -m pip install pydantic --no-cache-dir --force-reinstall
|
||||
"%PYTHON_EXE%" -m pip install typer --no-cache-dir --force-reinstall
|
||||
"%PYTHON_EXE%" -m pip install einops --no-cache-dir --force-reinstall
|
||||
"%PYTHON_EXE%" -m pip install omegaconf --no-cache-dir --force-reinstall
|
||||
"%PYTHON_EXE%" -m pip install easydict --no-cache-dir --force-reinstall
|
||||
"%PYTHON_EXE%" -m pip install yacs --no-cache-dir --force-reinstall
|
||||
)
|
||||
|
||||
"%PYTHON_EXE%" -c "import pydantic; import typer; import einops; import omegaconf; import easydict; import yacs; print('OK')" >nul 2>&1
|
||||
if errorlevel 1 (
|
||||
echo [X] Could not install iopaint dependencies
|
||||
echo Please try running: pip install pydantic typer einops omegaconf easydict yacs
|
||||
pause
|
||||
exit /b 1
|
||||
)
|
||||
)
|
||||
echo [OK] iopaint dependencies installed and verified
|
||||
|
||||
:: Download LaMA model
|
||||
echo.
|
||||
echo [*] Downloading LaMA model (196MB)...
|
||||
set LAMA_DIR=%USERPROFILE%\.cache\torch\hub\checkpoints
|
||||
set LAMA_FILE=%LAMA_DIR%\big-lama.pt
|
||||
|
||||
if not exist "%LAMA_FILE%" (
|
||||
if not exist "%LAMA_DIR%" mkdir "%LAMA_DIR%"
|
||||
powershell -Command "Invoke-WebRequest -Uri 'https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt' -OutFile '%LAMA_FILE%' -UseBasicParsing"
|
||||
if exist "%LAMA_FILE%" (
|
||||
echo [OK] LaMA model ready
|
||||
) else (
|
||||
echo [!] Warning: Could not download LaMA model
|
||||
echo It will be downloaded on first use
|
||||
)
|
||||
) else (
|
||||
echo [OK] LaMA model already exists
|
||||
echo [X] Could not prepare verified LaMA weights. Fix the error above and retry.
|
||||
pause
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
:: Download Florence-2 model
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
# WatermarkRemover-AI Setup Script
|
||||
$Host.UI.RawUI.WindowTitle = "WatermarkRemover-AI Setup"
|
||||
|
||||
Set-Location -LiteralPath $PSScriptRoot
|
||||
|
||||
$PYTHON_VERSION = "3.12.7"
|
||||
$PYTHON_DIR = "python"
|
||||
$PYTHON_EXE = "$PYTHON_DIR\python.exe"
|
||||
@@ -104,6 +106,15 @@ else {
|
||||
Write-Host " [OK] Python found" -ForegroundColor Green
|
||||
}
|
||||
|
||||
# Make local packages importable in fresh and existing embedded runtimes.
|
||||
$pthFile = Join-Path $PYTHON_DIR "python312._pth"
|
||||
if (Test-Path -LiteralPath $pthFile) {
|
||||
$pthLines = @(Get-Content -LiteralPath $pthFile)
|
||||
if ($pthLines -notcontains "..") {
|
||||
Add-Content -LiteralPath $pthFile -Value "`n.."
|
||||
}
|
||||
}
|
||||
|
||||
Write-Host ""
|
||||
Write-Host " [*] Installing dependencies..." -ForegroundColor Cyan
|
||||
Write-Host " This takes 5-10 minutes. Chill and learn something!" -ForegroundColor Magenta
|
||||
@@ -118,14 +129,16 @@ if ($CHINA_MODE) {
|
||||
& $PYTHON_EXE -m pip install --upgrade pip setuptools wheel 2>&1 | Out-Null
|
||||
}
|
||||
|
||||
# Install base deps with tips (legacy resolver to ignore conflicts)
|
||||
# Install dependencies with the normal resolver and retain the installation tips.
|
||||
if ($CHINA_MODE) {
|
||||
# For China: use Tsinghua mirror, skip PyTorch extra-index-url from requirements.txt
|
||||
$process = Start-Process -FilePath $PYTHON_EXE -ArgumentList "-m", "pip", "install", "--upgrade", "-r", "requirements.txt", "--no-cache-dir", "--use-deprecated=legacy-resolver", "-i", $PIP_INDEX_URL, "--trusted-host", "pypi.tuna.tsinghua.edu.cn" -NoNewWindow -PassThru
|
||||
# Use the selected package mirror.
|
||||
$process = Start-Process -FilePath $PYTHON_EXE -ArgumentList "-m", "pip", "install", "--upgrade", "-r", "requirements.txt", "--no-cache-dir", "-i", $PIP_INDEX_URL, "--trusted-host", "pypi.tuna.tsinghua.edu.cn" -NoNewWindow -PassThru
|
||||
} else {
|
||||
$process = Start-Process -FilePath $PYTHON_EXE -ArgumentList "-m", "pip", "install", "--upgrade", "-r", "requirements.txt", "--no-cache-dir", "--use-deprecated=legacy-resolver" -NoNewWindow -PassThru
|
||||
$process = Start-Process -FilePath $PYTHON_EXE -ArgumentList "-m", "pip", "install", "--upgrade", "-r", "requirements.txt", "--no-cache-dir" -NoNewWindow -PassThru
|
||||
}
|
||||
|
||||
# Retain the native process handle so ExitCode remains available after exit.
|
||||
$processHandle = $process.Handle
|
||||
$lastTipTime = Get-Date
|
||||
$currentTip = Get-Random -Maximum $tips.Count
|
||||
|
||||
@@ -145,138 +158,32 @@ while (-not $process.HasExited) {
|
||||
|
||||
Write-Host "`r "
|
||||
|
||||
# Legacy resolver can return non-zero even on success, so verify key packages
|
||||
$verifyResult = & $PYTHON_EXE -c "import torch; import transformers; import webview; import cv2; print('OK')" 2>&1
|
||||
if ($verifyResult -ne "OK") {
|
||||
if ($process.ExitCode -ne 0) {
|
||||
Write-Host ""
|
||||
Write-Host " [X] Failed to install dependencies" -ForegroundColor Red
|
||||
Read-Host " Press Enter to exit"
|
||||
exit 1
|
||||
}
|
||||
}
|
||||
|
||||
# Install iopaint separately without pulling its deps (we already have ours)
|
||||
Write-Host " [*] Installing iopaint (no deps)..." -ForegroundColor Cyan
|
||||
if ($CHINA_MODE) {
|
||||
$iopaintProcess = Start-Process -FilePath $PYTHON_EXE -ArgumentList "-m", "pip", "install", "--upgrade", "iopaint", "--no-deps", "--no-cache-dir", "-i", $PIP_INDEX_URL, "--trusted-host", "pypi.tuna.tsinghua.edu.cn" -NoNewWindow -PassThru
|
||||
} else {
|
||||
$iopaintProcess = Start-Process -FilePath $PYTHON_EXE -ArgumentList "-m", "pip", "install", "--upgrade", "iopaint", "--no-deps", "--no-cache-dir" -NoNewWindow -PassThru
|
||||
}
|
||||
$iopaintProcess.WaitForExit()
|
||||
|
||||
if ($iopaintProcess.ExitCode -ne 0) {
|
||||
Write-Host ""
|
||||
Write-Host " [X] Failed to install iopaint" -ForegroundColor Red
|
||||
Read-Host " Press Enter to exit"
|
||||
$process.WaitForExit()
|
||||
if ($process.ExitCode -ne 0) {
|
||||
Write-Host " [X] Failed to install dependencies" -ForegroundColor Red
|
||||
exit 1
|
||||
}
|
||||
Write-Host " [OK] iopaint installed" -ForegroundColor Green
|
||||
|
||||
# Install iopaint's required dependencies manually (subset needed for LaMA inpainting)
|
||||
Write-Host " [*] Installing iopaint dependencies..." -ForegroundColor Cyan
|
||||
if ($CHINA_MODE) {
|
||||
$iopaintDepsProcess = Start-Process -FilePath $PYTHON_EXE -ArgumentList "-m", "pip", "install", "pydantic", "typer", "einops", "omegaconf", "easydict", "yacs", "--no-cache-dir", "-i", $PIP_INDEX_URL, "--trusted-host", "pypi.tuna.tsinghua.edu.cn" -NoNewWindow -PassThru
|
||||
} else {
|
||||
$iopaintDepsProcess = Start-Process -FilePath $PYTHON_EXE -ArgumentList "-m", "pip", "install", "pydantic", "typer", "einops", "omegaconf", "easydict", "yacs", "--no-cache-dir" -NoNewWindow -PassThru
|
||||
}
|
||||
$iopaintDepsProcess.WaitForExit()
|
||||
|
||||
if ($iopaintDepsProcess.ExitCode -ne 0) {
|
||||
Write-Host ""
|
||||
Write-Host " [X] Failed to install iopaint dependencies" -ForegroundColor Red
|
||||
Read-Host " Press Enter to exit"
|
||||
& $PYTHON_EXE -m pip check
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host " [X] Dependency verification failed. Use a fresh application environment." -ForegroundColor Red
|
||||
exit 1
|
||||
}
|
||||
|
||||
# Verify iopaint dependencies are properly installed
|
||||
Write-Host " [*] Verifying iopaint dependencies..." -ForegroundColor Cyan
|
||||
$verifyIopaint = & $PYTHON_EXE -c "import pydantic; import typer; import einops; import omegaconf; import easydict; import yacs; print('OK')" 2>&1
|
||||
if ($verifyIopaint -ne "OK") {
|
||||
Write-Host ""
|
||||
Write-Host " [X] iopaint dependencies verification failed" -ForegroundColor Red
|
||||
Write-Host " Missing modules detected. Attempting reinstall..." -ForegroundColor Yellow
|
||||
|
||||
# Try installing one by one to identify issues
|
||||
$deps = @("pydantic", "typer", "einops", "omegaconf", "easydict", "yacs")
|
||||
foreach ($dep in $deps) {
|
||||
if ($CHINA_MODE) {
|
||||
& $PYTHON_EXE -m pip install $dep --no-cache-dir -i $PIP_INDEX_URL --trusted-host pypi.tuna.tsinghua.edu.cn 2>&1 | Out-Null
|
||||
} else {
|
||||
& $PYTHON_EXE -m pip install $dep --no-cache-dir 2>&1 | Out-Null
|
||||
}
|
||||
}
|
||||
|
||||
# Verify again
|
||||
$verifyAgain = & $PYTHON_EXE -c "import pydantic; import typer; import einops; import omegaconf; import easydict; import yacs; print('OK')" 2>&1
|
||||
if ($verifyAgain -ne "OK") {
|
||||
Write-Host " [X] Could not install iopaint dependencies" -ForegroundColor Red
|
||||
Write-Host " Please try running: pip install pydantic typer einops omegaconf easydict yacs" -ForegroundColor Yellow
|
||||
Read-Host " Press Enter to exit"
|
||||
exit 1
|
||||
}
|
||||
& $PYTHON_EXE -c "import remwm; import webview; import yaml; import psutil"
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host " [X] Failed to import application packages" -ForegroundColor Red
|
||||
exit 1
|
||||
}
|
||||
|
||||
Write-Host " [OK] Dependencies installed and verified" -ForegroundColor Green
|
||||
|
||||
# Download LaMA model directly from GitHub (avoids iopaint CLI dependency on fastapi)
|
||||
Write-Host ""
|
||||
Write-Host " [*] Downloading AI model (196MB)..." -ForegroundColor Cyan
|
||||
Write-Host ""
|
||||
Write-Host " Did you know?" -ForegroundColor DarkGray
|
||||
Write-Host ""
|
||||
|
||||
$lamaDir = Join-Path $env:USERPROFILE ".cache\torch\hub\checkpoints"
|
||||
$lamaFile = Join-Path $lamaDir "big-lama.pt"
|
||||
$lamaUrl = "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt"
|
||||
|
||||
if (-not (Test-Path $lamaFile)) {
|
||||
# Create directory if needed
|
||||
if (-not (Test-Path $lamaDir)) {
|
||||
New-Item -ItemType Directory -Path $lamaDir -Force | Out-Null
|
||||
}
|
||||
|
||||
try {
|
||||
# Show tips while downloading
|
||||
$job = Start-Job -ScriptBlock {
|
||||
param($url, $dest)
|
||||
Invoke-WebRequest -Uri $url -OutFile $dest -UseBasicParsing
|
||||
} -ArgumentList $lamaUrl, $lamaFile
|
||||
|
||||
$lastTipTime = Get-Date
|
||||
while ($job.State -eq "Running") {
|
||||
$now = Get-Date
|
||||
if (($now - $lastTipTime).TotalSeconds -ge 5) {
|
||||
$tip = $tips[$currentTip]
|
||||
$line = " $($tip.icon) $($tip.text)"
|
||||
$line = $line.PadRight(90)
|
||||
Write-Host "`r$line" -ForegroundColor $tip.color -NoNewline
|
||||
|
||||
$currentTip = ($currentTip + 1) % $tips.Count
|
||||
$lastTipTime = $now
|
||||
}
|
||||
Start-Sleep -Milliseconds 300
|
||||
}
|
||||
|
||||
Write-Host "`r "
|
||||
|
||||
$result = Receive-Job -Job $job
|
||||
Remove-Job -Job $job
|
||||
|
||||
if (Test-Path $lamaFile) {
|
||||
Write-Host " [OK] LaMA model ready" -ForegroundColor Green
|
||||
} else {
|
||||
Write-Host " [!] Warning: Could not download LaMA model" -ForegroundColor Yellow
|
||||
Write-Host " It will be downloaded on first use" -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
catch {
|
||||
Write-Host " [!] Warning: Could not download LaMA model" -ForegroundColor Yellow
|
||||
Write-Host " It will be downloaded on first use" -ForegroundColor Yellow
|
||||
}
|
||||
# All installers use the same checksum-verified, atomic model download.
|
||||
Write-Host " [*] Preparing LaMA model (196MB)..." -ForegroundColor Cyan
|
||||
if ($CHINA_MODE) {
|
||||
Write-Host " If GitHub is blocked, preseed the verified cache: docs/lama-runtime.md#restricted-networks" -ForegroundColor DarkGray
|
||||
}
|
||||
else {
|
||||
Write-Host " [OK] LaMA model already exists" -ForegroundColor Green
|
||||
& $PYTHON_EXE -m lama_inpaint download
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host " [X] Could not prepare verified LaMA weights. Fix the error above and retry." -ForegroundColor Red
|
||||
exit 1
|
||||
}
|
||||
|
||||
# Download Florence-2 model for watermark detection
|
||||
@@ -306,6 +213,7 @@ print('FLORENCE_OK')
|
||||
}
|
||||
|
||||
$florenceProcess = Start-Process -FilePath $PYTHON_EXE -ArgumentList "-c", "`"$florenceScript`"" -NoNewWindow -PassThru
|
||||
$florenceHandle = $florenceProcess.Handle
|
||||
|
||||
$lastTipTime = Get-Date
|
||||
while (-not $florenceProcess.HasExited) {
|
||||
@@ -324,6 +232,7 @@ while (-not $florenceProcess.HasExited) {
|
||||
|
||||
Write-Host "`r "
|
||||
|
||||
$florenceProcess.WaitForExit()
|
||||
if ($florenceProcess.ExitCode -ne 0) {
|
||||
Write-Host " [!] Warning: Could not download Florence-2 model" -ForegroundColor Yellow
|
||||
Write-Host " It will be downloaded on first use" -ForegroundColor Yellow
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
#!/usr/bin/env bash
|
||||
set -e
|
||||
cd -- "$(dirname -- "${BASH_SOURCE[0]}")"
|
||||
|
||||
echo ""
|
||||
echo " ============================================="
|
||||
@@ -83,74 +84,37 @@ echo " [*] Installing PyTorch..."
|
||||
if [ "$OS_TYPE" == "macos" ]; then
|
||||
# macOS: Install from main PyPI (supports MPS on Apple Silicon)
|
||||
if [ "$CHINA_MODE" == "1" ]; then
|
||||
pip install torch>=2.4.0 torchvision>=0.19.0 --no-cache-dir $PIP_MIRROR -q
|
||||
pip install "torch>=2.4.0" "torchvision>=0.19.0" --no-cache-dir $PIP_MIRROR -q
|
||||
else
|
||||
pip install torch>=2.4.0 torchvision>=0.19.0 --no-cache-dir -q
|
||||
pip install "torch>=2.4.0" "torchvision>=0.19.0" --no-cache-dir -q
|
||||
fi
|
||||
echo " [OK] PyTorch installed (MPS support on Apple Silicon)"
|
||||
else
|
||||
# Linux: Try CUDA first, fallback to CPU
|
||||
# Select one wheel index; mixing PyPI/mirrors can silently choose a CPU build.
|
||||
if command -v nvidia-smi &> /dev/null; then
|
||||
echo " [*] NVIDIA GPU detected, installing CUDA version..."
|
||||
if [ "$CHINA_MODE" == "1" ]; then
|
||||
pip install torch>=2.4.0 torchvision>=0.19.0 --extra-index-url https://download.pytorch.org/whl/cu124 --no-cache-dir $PIP_MIRROR -q
|
||||
else
|
||||
pip install torch>=2.4.0 torchvision>=0.19.0 --extra-index-url https://download.pytorch.org/whl/cu124 --no-cache-dir -q
|
||||
fi
|
||||
python -m pip --isolated install "torch>=2.4.0" "torchvision>=0.19.0" --index-url https://download.pytorch.org/whl/cu124 --no-cache-dir -q
|
||||
echo " [OK] PyTorch installed (CUDA 12.4)"
|
||||
else
|
||||
echo " [*] No NVIDIA GPU detected, installing CPU version..."
|
||||
if [ "$CHINA_MODE" == "1" ]; then
|
||||
pip install torch>=2.4.0 torchvision>=0.19.0 --no-cache-dir $PIP_MIRROR -q
|
||||
else
|
||||
pip install torch>=2.4.0 torchvision>=0.19.0 --no-cache-dir -q
|
||||
fi
|
||||
python -m pip --isolated install "torch>=2.4.0" "torchvision>=0.19.0" --index-url https://download.pytorch.org/whl/cpu --no-cache-dir -q
|
||||
echo " [OK] PyTorch installed (CPU)"
|
||||
fi
|
||||
fi
|
||||
|
||||
# Install other dependencies (without torch lines)
|
||||
echo " [*] Installing other dependencies..."
|
||||
if [ "$CHINA_MODE" == "1" ]; then
|
||||
pip install transformers>=4.50.0 diffusers>=0.30.0 "numpy<2" --no-cache-dir $PIP_MIRROR -q
|
||||
pip install "opencv-python-headless>=4.8.0,<4.12.0" "Pillow>=10.0.0" --no-cache-dir $PIP_MIRROR -q
|
||||
pip install pywebview>=4.0 --no-cache-dir $PIP_MIRROR -q
|
||||
pip install loguru click tqdm psutil pyyaml --no-cache-dir $PIP_MIRROR -q
|
||||
else
|
||||
pip install transformers>=4.50.0 diffusers>=0.30.0 "numpy<2" --no-cache-dir -q
|
||||
pip install "opencv-python-headless>=4.8.0,<4.12.0" "Pillow>=10.0.0" --no-cache-dir -q
|
||||
pip install pywebview>=4.0 --no-cache-dir -q
|
||||
pip install loguru click tqdm psutil pyyaml --no-cache-dir -q
|
||||
fi
|
||||
# Use the same dependency manifest as the Windows installers and manual installs.
|
||||
echo " [*] Installing application dependencies..."
|
||||
python -m pip install -r requirements.txt --no-cache-dir $PIP_MIRROR
|
||||
python -m pip check
|
||||
python -c "import remwm; import webview; import yaml; import psutil"
|
||||
echo " [OK] Dependencies installed and verified"
|
||||
|
||||
# Install iopaint separately (no deps to avoid conflicts)
|
||||
echo " [*] Installing iopaint..."
|
||||
# Shared checksum-verified, atomic download; failure must not look like success.
|
||||
echo " [*] Preparing LaMA model (~196MB)..."
|
||||
if [ "$CHINA_MODE" == "1" ]; then
|
||||
pip install iopaint --no-deps --no-cache-dir $PIP_MIRROR -q
|
||||
else
|
||||
pip install iopaint --no-deps --no-cache-dir -q
|
||||
fi
|
||||
|
||||
# Install iopaint's required dependencies manually (subset needed for LaMA inpainting)
|
||||
echo " [*] Installing iopaint dependencies..."
|
||||
if [ "$CHINA_MODE" == "1" ]; then
|
||||
pip install pydantic typer einops omegaconf easydict yacs --no-cache-dir $PIP_MIRROR -q
|
||||
else
|
||||
pip install pydantic typer einops omegaconf easydict yacs --no-cache-dir -q
|
||||
fi
|
||||
echo " [OK] Dependencies installed"
|
||||
|
||||
# Download LaMA model directly from GitHub (avoids iopaint CLI dependency on fastapi)
|
||||
echo " [*] Downloading LaMA model (~196MB)..."
|
||||
LAMA_DIR="$HOME/.cache/torch/hub/checkpoints"
|
||||
LAMA_FILE="$LAMA_DIR/big-lama.pt"
|
||||
if [ ! -f "$LAMA_FILE" ]; then
|
||||
mkdir -p "$LAMA_DIR"
|
||||
curl -L -o "$LAMA_FILE" "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt" || echo " [!] LaMA download failed, will retry on first use"
|
||||
echo " [OK] LaMA model downloaded"
|
||||
else
|
||||
echo " [OK] LaMA model already exists"
|
||||
echo " If GitHub is blocked, preseed the verified cache: docs/lama-runtime.md#restricted-networks"
|
||||
fi
|
||||
python -m lama_inpaint download
|
||||
|
||||
# Download Florence-2 model
|
||||
echo " [*] Downloading Florence-2 model (~1.5GB)..."
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
"""Exercise image and video application paths without detection/model downloads."""
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("torch", reason="requires the application runtime test job")
|
||||
import remwm
|
||||
from lama_inpaint.processing import InpaintProcessor
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def controlled_inpainting(monkeypatch):
|
||||
"""Keep detection and prediction deterministic and select the no-audio video path."""
|
||||
|
||||
def mask(image, *args, **kwargs):
|
||||
"""Select a fixed rectangle in the supplied image coordinates."""
|
||||
values = np.zeros((image.height, image.width), np.uint8)
|
||||
values[16:32, 16:32] = 255
|
||||
return Image.fromarray(values)
|
||||
|
||||
monkeypatch.setattr(remwm, "get_watermark_mask", mask)
|
||||
|
||||
def no_ffmpeg(*args, **kwargs):
|
||||
"""Choose the no-audio fallback deliberately for the silent fixture."""
|
||||
raise FileNotFoundError("controlled no-audio fixture")
|
||||
|
||||
monkeypatch.setattr(remwm.subprocess, "check_output", no_ffmpeg)
|
||||
return InpaintProcessor(lambda image, mask: np.full_like(image, (255, 0, 0)))
|
||||
|
||||
|
||||
def test_image_application_path_preserves_source_and_colors(
|
||||
tmp_path, controlled_inpainting
|
||||
):
|
||||
"""Check the real image routing path, source preservation and RGB/BGR conversion."""
|
||||
source, output = tmp_path / "source.png", tmp_path / "output.png"
|
||||
Image.new("RGB", (64, 48), (20, 60, 100)).save(source)
|
||||
before = source.read_bytes()
|
||||
remwm.handle_one(
|
||||
source,
|
||||
output,
|
||||
None,
|
||||
None,
|
||||
controlled_inpainting,
|
||||
"cpu",
|
||||
False,
|
||||
10,
|
||||
"PNG",
|
||||
False,
|
||||
)
|
||||
assert source.read_bytes() == before
|
||||
with Image.open(output) as image:
|
||||
assert image.size == (64, 48)
|
||||
assert image.getpixel((20, 20)) == (255, 0, 0)
|
||||
assert image.getpixel((0, 0)) == (20, 60, 100)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("detection_skip", [1, 2])
|
||||
def test_short_video_paths_keep_frames_and_dimensions(
|
||||
tmp_path, controlled_inpainting, monkeypatch, detection_skip
|
||||
):
|
||||
"""Exercise both video routes and check frame count, geometry and channel order."""
|
||||
source, output = tmp_path / "source.avi", tmp_path / "output.mp4"
|
||||
writer = cv2.VideoWriter(str(source), cv2.VideoWriter_fourcc(*"MJPG"), 5, (64, 48))
|
||||
assert writer.isOpened(), "test requires the OpenCV MJPG video writer"
|
||||
try:
|
||||
for _ in range(3):
|
||||
writer.write(np.full((48, 64, 3), (100, 60, 20), np.uint8))
|
||||
finally:
|
||||
writer.release()
|
||||
# Two-pass mode caches detected boxes, while single-pass mode uses the mask.
|
||||
monkeypatch.setattr(
|
||||
remwm,
|
||||
"detect_only",
|
||||
lambda *a, **k: [
|
||||
{"bbox": [16, 16, 31, 31], "area_percent": 8, "accepted": True}
|
||||
],
|
||||
)
|
||||
before = source.read_bytes()
|
||||
remwm.handle_one(
|
||||
source,
|
||||
output,
|
||||
None,
|
||||
None,
|
||||
controlled_inpainting,
|
||||
"cpu",
|
||||
False,
|
||||
10,
|
||||
"MP4",
|
||||
False,
|
||||
detection_skip=detection_skip,
|
||||
)
|
||||
assert source.read_bytes() == before
|
||||
capture = cv2.VideoCapture(str(output))
|
||||
frames = []
|
||||
try:
|
||||
while True:
|
||||
ok, frame = capture.read()
|
||||
if not ok:
|
||||
break
|
||||
frames.append(frame)
|
||||
finally:
|
||||
capture.release()
|
||||
assert len(frames) == 3
|
||||
for frame in frames:
|
||||
assert frame.shape == (48, 64, 3)
|
||||
# Lossy video encoding needs a tolerance; the channel order must survive.
|
||||
assert frame[23, 23, 2] > 200
|
||||
assert frame[23, 23, 0] < 30
|
||||
@@ -0,0 +1,146 @@
|
||||
"""Download integrity and failure behavior using tiny local byte streams."""
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import hashlib
|
||||
import io
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from lama_inpaint import model_store as store
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def artifact(monkeypatch):
|
||||
"""Pin verification to tiny controlled bytes so download tests need no real weights."""
|
||||
data = b"controlled model fixture"
|
||||
monkeypatch.setattr(store, "MODEL_SIZE", len(data))
|
||||
monkeypatch.setattr(store, "MODEL_SHA256", hashlib.sha256(data).hexdigest())
|
||||
return data
|
||||
|
||||
|
||||
def test_download_then_offline_cache_reuse(tmp_path, monkeypatch, artifact):
|
||||
"""Download once and verify that both online and offline reuse avoid another request."""
|
||||
calls = []
|
||||
|
||||
def response(url, timeout):
|
||||
"""Record a network attempt and supply the controlled artifact bytes."""
|
||||
assert url == store.MODEL_URL and timeout == 30
|
||||
calls.append(url)
|
||||
return io.BytesIO(artifact)
|
||||
|
||||
monkeypatch.setattr(store, "urlopen", response)
|
||||
target = store.ensure_model(tmp_path)
|
||||
assert target.read_bytes() == artifact
|
||||
assert store.ensure_model(tmp_path, download=False) == target
|
||||
assert store.ensure_model(tmp_path) == target
|
||||
assert len(calls) == 1
|
||||
assert not list(tmp_path.glob("*.part"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [b"", b"truncated", b"x" * 200])
|
||||
def test_bad_download_never_published(tmp_path, monkeypatch, artifact, bad):
|
||||
"""Reject missing, truncated or oversized transfers without leaving a cache artifact."""
|
||||
monkeypatch.setattr(store, "urlopen", lambda *a, **k: io.BytesIO(bad))
|
||||
with pytest.raises(store.ModelError):
|
||||
store.ensure_model(tmp_path)
|
||||
assert not (tmp_path / store.MODEL_NAME).exists()
|
||||
assert not list(tmp_path.glob("*.part"))
|
||||
|
||||
|
||||
def test_same_size_tampering_rejected_offline(tmp_path, artifact):
|
||||
"""Require the digest as well as size when accepting offline weights."""
|
||||
(tmp_path / store.MODEL_NAME).write_bytes(b"x" * len(artifact))
|
||||
with pytest.raises(store.ModelError, match="No verified"):
|
||||
store.ensure_model(tmp_path, download=False)
|
||||
|
||||
|
||||
def test_missing_offline_model_does_not_connect(tmp_path, monkeypatch, artifact):
|
||||
"""Fail an absent offline cache without opening a network connection."""
|
||||
|
||||
def fail(*a, **k):
|
||||
"""Reject an unexpected network request in an offline test."""
|
||||
pytest.fail("offline mode must not connect")
|
||||
|
||||
monkeypatch.setattr(store, "urlopen", fail)
|
||||
with pytest.raises(store.ModelError):
|
||||
store.ensure_model(tmp_path, download=False)
|
||||
|
||||
|
||||
def test_interrupted_download_preserves_old_file(tmp_path, monkeypatch, artifact):
|
||||
"""Keep the old artifact and remove temporary data after a connection interruption."""
|
||||
old = tmp_path / store.MODEL_NAME
|
||||
old.write_bytes(b"old invalid file")
|
||||
|
||||
class Interrupted(io.BytesIO):
|
||||
def read(self, size=-1):
|
||||
"""Supply one partial chunk, then simulate an interrupted connection."""
|
||||
if self.tell():
|
||||
raise OSError("connection interrupted")
|
||||
return super().read(5)
|
||||
|
||||
monkeypatch.setattr(store, "urlopen", lambda *a, **k: Interrupted(artifact))
|
||||
with pytest.raises(store.ModelError, match="interrupted"):
|
||||
store.ensure_model(tmp_path)
|
||||
assert old.read_bytes() == b"old invalid file"
|
||||
assert not list(tmp_path.glob("*.part"))
|
||||
|
||||
|
||||
def test_corrupt_cache_replaced_only_after_verification(
|
||||
tmp_path, monkeypatch, artifact
|
||||
):
|
||||
"""Replace an invalid existing artifact only with the complete verified fixture."""
|
||||
target = tmp_path / store.MODEL_NAME
|
||||
target.write_bytes(b"bad")
|
||||
monkeypatch.setattr(store, "urlopen", lambda *a, **k: io.BytesIO(artifact))
|
||||
assert store.ensure_model(tmp_path).read_bytes() == artifact
|
||||
|
||||
|
||||
def test_concurrent_downloads_share_verified_artifact(tmp_path, monkeypatch, artifact):
|
||||
"""Concurrent callers must obtain one verified download through the artifact lock."""
|
||||
calls = []
|
||||
|
||||
def response(*a, **k):
|
||||
"""Record a network attempt and supply the controlled artifact bytes."""
|
||||
calls.append(1)
|
||||
return io.BytesIO(artifact)
|
||||
|
||||
monkeypatch.setattr(store, "urlopen", response)
|
||||
with ThreadPoolExecutor(max_workers=4) as pool:
|
||||
paths = list(pool.map(lambda _: store.ensure_model(tmp_path), range(4)))
|
||||
assert all(p.read_bytes() == artifact for p in paths)
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
def test_torch_cache_environment(monkeypatch, tmp_path):
|
||||
"""Apply Torch cache precedence before falling back to the XDG cache directory."""
|
||||
monkeypatch.setenv("TORCH_HOME", str(tmp_path / "custom"))
|
||||
assert store.default_cache_dir() == tmp_path / "custom" / "hub" / "checkpoints"
|
||||
monkeypatch.delenv("TORCH_HOME")
|
||||
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path))
|
||||
assert store.default_cache_dir() == tmp_path / "torch" / "hub" / "checkpoints"
|
||||
|
||||
|
||||
def test_verification_rewinds_stream(artifact):
|
||||
"""Leave a verified open stream ready for deserialization from its beginning."""
|
||||
stream = io.BytesIO(artifact)
|
||||
stream.seek(4)
|
||||
store.verify_stream(stream)
|
||||
assert stream.tell() == 0
|
||||
|
||||
|
||||
def test_incomplete_http_response_is_reported_and_cleaned(
|
||||
tmp_path, monkeypatch, artifact
|
||||
):
|
||||
"""Convert an incomplete HTTP response into a controlled error without residue."""
|
||||
from http.client import IncompleteRead
|
||||
|
||||
def interrupted(*args, **kwargs):
|
||||
"""Simulate an incomplete HTTP response before an artifact can be published."""
|
||||
raise IncompleteRead(b"partial", len(artifact))
|
||||
|
||||
monkeypatch.setattr(store, "urlopen", interrupted)
|
||||
with pytest.raises(store.ModelError, match="Could not download"):
|
||||
store.ensure_model(tmp_path)
|
||||
assert not (tmp_path / store.MODEL_NAME).exists()
|
||||
assert not list(tmp_path.glob("*.part"))
|
||||
@@ -0,0 +1,119 @@
|
||||
"""Observable crop/pad/composite behavior, with no model download."""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from lama_inpaint.processing import InpaintProcessor
|
||||
|
||||
|
||||
def paint_red(image, mask):
|
||||
"""Validate padded geometry and return an unmistakable RGB prediction."""
|
||||
assert image.shape[:2] == mask.shape
|
||||
assert image.shape[0] % 8 == image.shape[1] % 8 == 0
|
||||
return np.full_like(image, (255, 0, 0))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(1, 1), (7, 9), (16, 16), (801, 19), (17, 803)])
|
||||
def test_masked_pixels_change_and_originals_survive(shape):
|
||||
"""Check BGR output, edge masks, exact unmasked pixels and immutable source arrays."""
|
||||
image = np.full((*shape, 3), (13, 47, 91), dtype=np.uint8)
|
||||
mask = np.zeros(shape, dtype=np.uint8)
|
||||
mask[0, 0] = 255
|
||||
mask[-1, -1] = 255
|
||||
before = image.copy(), mask.copy()
|
||||
out = InpaintProcessor(paint_red)(image, mask)
|
||||
assert out.shape == image.shape and out.dtype == np.uint8
|
||||
np.testing.assert_array_equal(out[mask == 0], image[:, :, ::-1][mask == 0])
|
||||
assert np.all(out[mask == 255] == (0, 0, 255))
|
||||
np.testing.assert_array_equal(image, before[0])
|
||||
np.testing.assert_array_equal(mask, before[1])
|
||||
assert not np.shares_memory(out, image)
|
||||
|
||||
|
||||
def test_empty_mask_skips_model_and_returns_independent_bgr():
|
||||
"""Empty masks must skip prediction and return an independent channel-swapped copy."""
|
||||
image = np.full((801, 4, 3), (11, 29, 83), dtype=np.uint8)
|
||||
|
||||
def fail(*args):
|
||||
"""Fail immediately if an empty-mask request reaches the predictor."""
|
||||
pytest.fail("empty mask must not invoke the model")
|
||||
|
||||
out = InpaintProcessor(fail)(image, np.zeros((801, 4), np.uint8))
|
||||
np.testing.assert_array_equal(out, image[:, :, ::-1])
|
||||
assert not np.shares_memory(out, image)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(3, 5), (801, 3)])
|
||||
def test_full_mask_and_soft_mask(shape):
|
||||
"""Check full replacement and proportional blending above and below the crop trigger."""
|
||||
image = np.full((*shape, 3), (11, 29, 83), dtype=np.uint8)
|
||||
mask = np.full(shape, 64, dtype=np.uint8)
|
||||
out = InpaintProcessor(paint_red)(image, mask)
|
||||
expected = (
|
||||
np.array((0, 0, 255)) * (64 / 255) + np.array((83, 29, 11)) * (1 - 64 / 255)
|
||||
).astype(np.uint8)
|
||||
assert np.all(out == expected)
|
||||
assert np.all(
|
||||
InpaintProcessor(paint_red)(image, np.full(shape, 255, np.uint8)) == (0, 0, 255)
|
||||
)
|
||||
|
||||
|
||||
def test_overlapping_crops_always_use_original_context():
|
||||
"""Overlapping context must not feed an earlier prediction into another model call."""
|
||||
image = np.full((900, 900, 3), 37, np.uint8)
|
||||
mask = np.zeros((900, 900), np.uint8)
|
||||
mask[200:205, 200:205] = 255
|
||||
mask[220:225, 220:225] = 255
|
||||
calls = []
|
||||
|
||||
def predictor(rgb, selected):
|
||||
"""Inspect the padded source crop and return a controlled prediction."""
|
||||
assert np.all(rgb == 37)
|
||||
calls.append(rgb.shape)
|
||||
return paint_red(rgb, selected)
|
||||
|
||||
out = InpaintProcessor(predictor)(image, mask)
|
||||
assert len(calls) == 2
|
||||
assert np.all(out[mask == 0] == 37)
|
||||
assert np.all(out[mask != 0] == (0, 0, 255))
|
||||
|
||||
|
||||
def test_small_image_symmetric_padding():
|
||||
"""Extend tiny inputs symmetrically without changing their original-size result."""
|
||||
image = np.arange(18, dtype=np.uint8).reshape(2, 3, 3)
|
||||
seen = []
|
||||
|
||||
def predictor(rgb, mask):
|
||||
"""Inspect the padded source crop and return a controlled prediction."""
|
||||
seen.append(rgb.copy())
|
||||
return rgb
|
||||
|
||||
out = InpaintProcessor(predictor)(image, np.full((2, 3), 255, np.uint8))
|
||||
np.testing.assert_array_equal(out, image[:, :, ::-1])
|
||||
assert seen[0].shape == (32, 32, 3)
|
||||
np.testing.assert_array_equal(seen[0][0, 3], image[0, 2])
|
||||
np.testing.assert_array_equal(seen[0][2, 0], image[1, 0])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"image,mask",
|
||||
[
|
||||
(np.zeros((0, 2, 3), np.uint8), np.zeros((0, 2), np.uint8)),
|
||||
(np.zeros((2, 2, 4), np.uint8), np.zeros((2, 2), np.uint8)),
|
||||
(np.zeros((2, 2, 3), np.float32), np.zeros((2, 2), np.uint8)),
|
||||
(np.zeros((2, 2, 3), np.uint8), np.zeros((2, 3), np.uint8)),
|
||||
(np.zeros((2, 2, 3), np.uint8), np.zeros((2, 2), bool)),
|
||||
],
|
||||
)
|
||||
def test_reject_invalid_inputs(image, mask):
|
||||
"""Reject invalid dimensions, channels, dtypes and mismatched masks before inference."""
|
||||
with pytest.raises(ValueError):
|
||||
InpaintProcessor(paint_red)(image, mask)
|
||||
|
||||
|
||||
def test_reject_invalid_model_output():
|
||||
"""Reject a predictor result that violates the padded uint8 RGB contract."""
|
||||
with pytest.raises(ValueError, match="predictor"):
|
||||
InpaintProcessor(lambda *_: np.zeros((8, 8, 3)))(
|
||||
np.zeros((2, 2, 3), np.uint8), np.full((2, 2), 255, np.uint8)
|
||||
)
|
||||
@@ -0,0 +1,69 @@
|
||||
"""Real TorchScript execution on a generated fixture; no downloaded models."""
|
||||
|
||||
import hashlib
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch", reason="requires the application runtime test job")
|
||||
from lama_inpaint import model_store
|
||||
from lama_inpaint.runtime import LamaInpaint
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def scripted_weights(tmp_path, monkeypatch):
|
||||
"""Create and checksum a tiny local TorchScript fixture without model downloads."""
|
||||
|
||||
class Paint(torch.nn.Module):
|
||||
def forward(self, image, mask):
|
||||
# Exercise both tensor inputs; model writes red into selected pixels.
|
||||
"""Use both tensor inputs to replace selected pixels with RGB red."""
|
||||
color = torch.zeros_like(image)
|
||||
color[:, 0] = 1
|
||||
return image * (1 - mask) + color * mask
|
||||
|
||||
traced = torch.jit.trace(
|
||||
Paint(), (torch.zeros(1, 3, 8, 8), torch.zeros(1, 1, 8, 8, dtype=torch.int64))
|
||||
)
|
||||
path = tmp_path / model_store.MODEL_NAME
|
||||
torch.jit.save(traced, path)
|
||||
data = path.read_bytes()
|
||||
monkeypatch.setattr(model_store, "MODEL_SIZE", len(data))
|
||||
monkeypatch.setattr(model_store, "MODEL_SHA256", hashlib.sha256(data).hexdigest())
|
||||
return tmp_path
|
||||
|
||||
|
||||
def test_cpu_scripted_model_end_to_end(scripted_weights):
|
||||
"""Exercise verified loading, tensor conversion and compositing with a real TorchScript call."""
|
||||
model = LamaInpaint("cpu", cache_dir=scripted_weights, download=False)
|
||||
image = np.full((13, 17, 3), (23, 51, 81), np.uint8)
|
||||
mask = np.zeros((13, 17), np.uint8)
|
||||
mask[2:7, 3:9] = 255
|
||||
result = model(image, mask)
|
||||
assert np.all(result[mask == 255] == (0, 0, 255))
|
||||
np.testing.assert_array_equal(result[mask == 0], image[:, :, ::-1][mask == 0])
|
||||
|
||||
|
||||
def test_corrupt_weights_are_never_deserialized(scripted_weights, monkeypatch):
|
||||
"""Reject modified artifact bytes before invoking the TorchScript loader."""
|
||||
(scripted_weights / model_store.MODEL_NAME).write_bytes(b"corrupt")
|
||||
monkeypatch.setattr(
|
||||
torch.jit,
|
||||
"load",
|
||||
lambda *a, **k: pytest.fail("must not load unverified weights"),
|
||||
)
|
||||
with pytest.raises(model_store.ModelError):
|
||||
LamaInpaint("cpu", cache_dir=scripted_weights, download=False)
|
||||
|
||||
|
||||
def test_reject_unvalidated_device_before_loading():
|
||||
"""Reject MPS explicitly rather than promising unvalidated operator support."""
|
||||
with pytest.raises(ValueError, match="MPS"):
|
||||
LamaInpaint("mps", download=False)
|
||||
|
||||
|
||||
def test_cuda_unavailable_is_explicit(monkeypatch):
|
||||
"""Report unavailable CUDA without silently selecting a different device."""
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
|
||||
with pytest.raises(RuntimeError, match="CUDA"):
|
||||
LamaInpaint("cuda", download=False)
|
||||
@@ -0,0 +1,97 @@
|
||||
"""Execute Linux setup with controlled command stubs; never install or download."""
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
sys.platform != "linux", reason="Linux installer branches"
|
||||
)
|
||||
REPO = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def run_setup(tmp_path, *, gpu, mirror, fail_install=False):
|
||||
"""Record selected indexes and stage ordering from the actual shell script."""
|
||||
project = tmp_path / "project with spaces"
|
||||
project.mkdir()
|
||||
shutil.copyfile(REPO / "setup.sh", project / "setup.sh")
|
||||
bindir = tmp_path / "bin"
|
||||
bindir.mkdir()
|
||||
program = tmp_path / "fake_python.py"
|
||||
program.write_text(
|
||||
r"""import json, os, sys
|
||||
from pathlib import Path
|
||||
args = sys.argv[1:]
|
||||
with open(os.environ['WMR_TEST_LOG'], 'a') as log:
|
||||
log.write(json.dumps(args) + '\n')
|
||||
if args[:1] == ['-c'] and 'sys.version_info' in args[1]:
|
||||
print('3.12')
|
||||
elif args[:2] == ['-m', 'venv']:
|
||||
activate = Path(args[2]) / 'bin' / 'activate'
|
||||
activate.parent.mkdir(parents=True)
|
||||
activate.write_text('# controlled environment\n')
|
||||
elif 'requirements.txt' in args and os.environ.get('WMR_TEST_FAIL') == '1':
|
||||
sys.exit(17)
|
||||
""",
|
||||
encoding="utf-8",
|
||||
)
|
||||
wrapper = '#!/bin/sh\nexec "$WMR_TEST_PYTHON" "$WMR_TEST_PROGRAM" "$@"\n'
|
||||
for name in ("python3.12", "python", "pip"):
|
||||
executable = bindir / name
|
||||
executable.write_text(wrapper, encoding="utf-8")
|
||||
executable.chmod(0o755)
|
||||
if gpu:
|
||||
executable = bindir / "nvidia-smi"
|
||||
executable.write_text("#!/bin/sh\nexit 0\n", encoding="utf-8")
|
||||
executable.chmod(0o755)
|
||||
log = tmp_path / "commands.jsonl"
|
||||
env = os.environ.copy()
|
||||
env.update(
|
||||
PATH=f"{bindir}:/usr/bin:/bin",
|
||||
WMR_TEST_PYTHON=sys.executable,
|
||||
WMR_TEST_PROGRAM=str(program),
|
||||
WMR_TEST_LOG=str(log),
|
||||
WMR_TEST_FAIL="1" if fail_install else "0",
|
||||
PIP_EXTRA_INDEX_URL="https://packages.invalid/simple",
|
||||
)
|
||||
result = subprocess.run(
|
||||
["/bin/bash", str(project / "setup.sh")],
|
||||
input=("y" if mirror else "n") + "n",
|
||||
env=env,
|
||||
text=True,
|
||||
capture_output=True,
|
||||
timeout=30,
|
||||
)
|
||||
commands = [json.loads(line) for line in log.read_text().splitlines()]
|
||||
return result, commands
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gpu", [False, True])
|
||||
@pytest.mark.parametrize("mirror", [False, True])
|
||||
def test_torch_backend_index_is_unambiguous(tmp_path, gpu, mirror):
|
||||
"""Verify CPU/CUDA index selection independently of the general package mirror."""
|
||||
result, commands = run_setup(tmp_path, gpu=gpu, mirror=mirror)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
torch_install = next(args for args in commands if "torch>=2.4.0" in args)
|
||||
expected = "https://download.pytorch.org/whl/" + ("cu124" if gpu else "cpu")
|
||||
assert torch_install[torch_install.index("--index-url") + 1] == expected
|
||||
assert "--isolated" in torch_install
|
||||
assert "--extra-index-url" not in torch_install
|
||||
assert "-i" not in torch_install
|
||||
app_install = next(args for args in commands if "requirements.txt" in args)
|
||||
assert ("-i" in app_install) == mirror
|
||||
assert ["-m", "pip", "check"] in commands
|
||||
assert ["-m", "lama_inpaint", "download"] in commands
|
||||
|
||||
|
||||
def test_install_failure_stops_before_model_preparation(tmp_path):
|
||||
"""A failed dependency install must stop before download or success reporting."""
|
||||
result, commands = run_setup(tmp_path, gpu=False, mirror=False, fail_install=True)
|
||||
assert result.returncode == 17
|
||||
assert ["-m", "lama_inpaint", "download"] not in commands
|
||||
assert "Setup complete!" not in result.stdout
|
||||
Reference in new issue
Block a user