From dc6c1a1ea571a13e974afd0ac76c63a3694fa827 Mon Sep 17 00:00:00 2001 From: Diego Lucio D'Onofrio Date: Wed, 16 Sep 2026 16:19:00 +0200 Subject: [PATCH] Add verbose resumable LaMA downloads --- lama_inpaint/__main__.py | 11 +- lama_inpaint/model_store.py | 218 +++++++++++++++++++++++++++++---- setup.ps1 | 2 +- setup.sh | 2 +- tests/test_lama_model_store.py | 27 +++- tests/test_setup_shell.py | 4 +- 6 files changed, 228 insertions(+), 36 deletions(-) diff --git a/lama_inpaint/__main__.py b/lama_inpaint/__main__.py index fc25b62..aa84dea 100644 --- a/lama_inpaint/__main__.py +++ b/lama_inpaint/__main__.py @@ -16,9 +16,18 @@ def main(): action="store_true", help="Verify cached weights without network access", ) + parser.add_argument( + "--verbose", + action="store_true", + help="Show download progress, resume information and verification details", + ) args = parser.parse_args() try: - path = ensure_model(args.cache_dir, download=not args.offline) + path = ensure_model( + args.cache_dir, + download=not args.offline, + verbose=args.verbose, + ) except (ModelError, OSError) as exc: print(f"LaMA model preparation failed: {exc}", file=sys.stderr) return 1 diff --git a/lama_inpaint/model_store.py b/lama_inpaint/model_store.py index 84239c9..b814fee 100644 --- a/lama_inpaint/model_store.py +++ b/lama_inpaint/model_store.py @@ -4,9 +4,9 @@ import hashlib from http.client import HTTPException import os from pathlib import Path -import tempfile +import sys import time -from urllib.request import urlopen +from urllib.request import Request, urlopen from filelock import FileLock, Timeout @@ -16,6 +16,10 @@ MODEL_URL = ( MODEL_SHA256 = "344c77bbcb158f17dd143070d1e789f38a66c04202311ae3a258ef66667a9ea9" MODEL_SIZE = 205669692 MODEL_NAME = "big-lama.pt" +DOWNLOAD_TIMEOUT = 30 +DOWNLOAD_DEADLINE = 600 +DOWNLOAD_CHUNK_SIZE = 1024 * 1024 +PROGRESS_INTERVAL = 1.0 class ModelError(RuntimeError): @@ -56,7 +60,7 @@ def _is_valid(path): return False -def ensure_model(cache_dir=None, *, download=True, progress=None): +def ensure_model(cache_dir=None, *, download=True, progress=None, verbose=False): """Return verified weights, downloading atomically when allowed. Downloads use a unique temporary file and a per-artifact cross-process lock. @@ -78,41 +82,201 @@ def ensure_model(cache_dir=None, *, download=True, progress=None): with FileLock(str(target) + ".lock", timeout=600): if _is_valid(target): return target - _download(target, progress=progress) + _download(target, progress=progress, verbose=verbose) except Timeout as exc: raise ModelError("Timed out waiting for another LaMA model download.") from exc return target -def _download(target, progress=None): - """Publish only a complete verified transfer; retain the old file on failure.""" - temporary = None +def _format_bytes(size): + """Format a byte count for human-readable progress output.""" + units = ("B", "KB", "MB", "GB") + value = float(size) + for unit in units: + if value < 1024 or unit == units[-1]: + return f"{value:.1f} {unit}" if unit != "B" else f"{int(value)} B" + value /= 1024 + + +def _format_duration(seconds): + """Format a duration without requiring a third-party progress package.""" + if seconds <= 0 or seconds == float("inf"): + return "--:--" + seconds = int(seconds) + minutes, seconds = divmod(seconds, 60) + hours, minutes = divmod(minutes, 60) + if hours: + return f"{hours:d}:{minutes:02d}:{seconds:02d}" + return f"{minutes:02d}:{seconds:02d}" + + +def _verbose(verbose, message): + """Write downloader diagnostics immediately when verbose mode is enabled.""" + if verbose: + print(f"[LaMA] {message}", file=sys.stderr, flush=True) + + +def _response_status(response): + """Read an HTTP status while remaining compatible with simple test streams.""" + status = getattr(response, "status", None) + if status is None: + getcode = getattr(response, "getcode", None) + status = getcode() if getcode is not None else None + return status + + +def _response_header(response, name): + """Read a response header from real HTTP responses or return None for test streams.""" + headers = getattr(response, "headers", None) + if headers is None: + return None + return headers.get(name) + + +def _download(target, progress=None, *, verbose=False): + """Download, resume and verify the model before publishing it atomically.""" + partial = target.with_name(f".{target.name}.part") + started_at = time.monotonic() + last_progress = 0.0 + 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: + existing = partial.stat().st_size if partial.exists() else 0 + if existing > MODEL_SIZE: + _verbose(verbose, f"Partial file exceeds expected size; restarting: {partial}") + partial.unlink() + existing = 0 + + if existing == MODEL_SIZE: + _verbose(verbose, "Partial file is complete; verifying SHA-256...") + if _is_valid(partial): + os.replace(partial, target) + _verbose(verbose, "Partial file verified and published as the model.") + return + _verbose(verbose, "Complete partial file failed SHA-256; restarting.") + partial.unlink() + existing = 0 + + _verbose( + verbose, + f"Target: {target} | Partial: {partial} | " + f"Expected size: {_format_bytes(MODEL_SIZE)}", + ) + + response = None + resumed = existing > 0 + if resumed: + _verbose( + verbose, + f"Resuming from {_format_bytes(existing)} " + f"({existing / MODEL_SIZE:.1%}) using HTTP Range...", + ) + request = Request(MODEL_URL, headers={"Range": f"bytes={existing}-"}) + response = urlopen(request, timeout=DOWNLOAD_TIMEOUT) + status = _response_status(response) + if status != 206: + _verbose( + verbose, + f"Server did not accept resume (HTTP {status or 'unknown'}); " + "restarting the transfer.", + ) + response.close() + response = None + partial.unlink() + existing = 0 + resumed = False + + if response is None: + _verbose(verbose, f"Downloading from {MODEL_URL}") + response = urlopen(MODEL_URL, timeout=DOWNLOAD_TIMEOUT) + + content_range = _response_header(response, "Content-Range") + if resumed and content_range: + expected_prefix = f"bytes {existing}-" + if not content_range.startswith(expected_prefix): + response.close() + raise ModelError( + f"Server returned an unexpected range ({content_range}); " + f"partial file retained at {partial} for retry." + ) + + mode = "ab" if existing else "wb" + downloaded = existing + previous_downloaded = downloaded + progress_started_at = time.monotonic() + last_progress = 0.0 + _verbose( + verbose, + f"Transfer {'resumed' if resumed else 'started'}; " + "interrupted transfers will retain their progress.", + ) + + with response: + with partial.open(mode) as output: + while True: + if time.monotonic() - started_at > DOWNLOAD_DEADLINE: raise ModelError( - "LaMA download exceeded its size or time limit." + f"Download deadline reached at {_format_bytes(downloaded)}/" + f"{_format_bytes(MODEL_SIZE)}; partial file retained at {partial}." ) + chunk = response.read(DOWNLOAD_CHUNK_SIZE) + if not chunk: + break output.write(chunk) + downloaded += len(chunk) + if downloaded > MODEL_SIZE: + raise ModelError( + f"Download exceeds expected size " + f"({_format_bytes(downloaded)} > {_format_bytes(MODEL_SIZE)}); " + f"partial file retained at {partial}." + ) + now = time.monotonic() + if verbose and ( + now - last_progress >= PROGRESS_INTERVAL + or downloaded >= MODEL_SIZE + ): + elapsed = max(now - progress_started_at, 0.001) + speed = (downloaded - previous_downloaded) / elapsed + remaining = max(MODEL_SIZE - downloaded, 0) + eta = remaining / speed if speed > 0 else float("inf") + print( + f"\r[LaMA] {_format_bytes(downloaded)}/" + f"{_format_bytes(MODEL_SIZE)} " + f"({downloaded / MODEL_SIZE:.1%}) | " + f"{_format_bytes(speed)}/s | ETA {_format_duration(eta)}", + end="", + file=sys.stderr, + flush=True, + ) + last_progress = now if progress: - progress(dict(status="downloading", model="LaMA", current=size, total=MODEL_SIZE)) - output.flush() - os.fsync(output.fileno()) - with temporary.open("rb") as stream: + progress( + dict( + status="downloading", + model="LaMA", + current=downloaded, + total=MODEL_SIZE, + ) + ) + output.flush() + os.fsync(output.fileno()) + + if verbose: + print(file=sys.stderr, flush=True) + if downloaded != MODEL_SIZE: + raise ModelError( + f"Incomplete download: {_format_bytes(downloaded)}/" + f"{_format_bytes(MODEL_SIZE)}; partial file retained at {partial}." + ) + + _verbose(verbose, "Download complete; verifying size and SHA-256...") + with partial.open("rb") as stream: verify_stream(stream) - os.replace(temporary, target) + os.replace(partial, target) + _verbose(verbose, f"Model verified and saved to {target}") + except ModelError: + raise except (OSError, ValueError, HTTPException) as exc: raise ModelError( - f"Could not download LaMA weights: {exc}. Retry python -m lama_inpaint download." + f"Could not download LaMA: {exc}. " + f"Partial file retained at {partial}; run the command again." ) from exc - finally: - if temporary is not None: - temporary.unlink(missing_ok=True) diff --git a/setup.ps1 b/setup.ps1 index 20cbd06..a7b442e 100644 --- a/setup.ps1 +++ b/setup.ps1 @@ -180,7 +180,7 @@ 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 } -& $PYTHON_EXE -m lama_inpaint download +& $PYTHON_EXE -m lama_inpaint download --verbose if ($LASTEXITCODE -ne 0) { Write-Host " [X] Could not prepare verified LaMA weights. Fix the error above and retry." -ForegroundColor Red exit 1 diff --git a/setup.sh b/setup.sh index 9333dca..1d1e7c9 100755 --- a/setup.sh +++ b/setup.sh @@ -114,7 +114,7 @@ echo " [*] Preparing LaMA model (~196MB)..." if [ "$CHINA_MODE" == "1" ]; then echo " If GitHub is blocked, preseed the verified cache: docs/lama-runtime.md#restricted-networks" fi -python -m lama_inpaint download +python -m lama_inpaint download --verbose # Download Florence-2 model echo " [*] Downloading Florence-2 model (~1.5GB)..." diff --git a/tests/test_lama_model_store.py b/tests/test_lama_model_store.py index 2ee73ae..2ee78e3 100644 --- a/tests/test_lama_model_store.py +++ b/tests/test_lama_model_store.py @@ -40,12 +40,12 @@ def test_download_then_offline_cache_reuse(tmp_path, monkeypatch, artifact): @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.""" + """Reject bad transfers without publishing them, while retaining resumable data.""" 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")) + assert (tmp_path / f".{store.MODEL_NAME}.part").exists() def test_same_size_tampering_rejected_offline(tmp_path, artifact): @@ -68,7 +68,7 @@ def test_missing_offline_model_does_not_connect(tmp_path, monkeypatch, artifact) def test_interrupted_download_preserves_old_file(tmp_path, monkeypatch, artifact): - """Keep the old artifact and remove temporary data after a connection interruption.""" + """Keep the old artifact and resumable data after a connection interruption.""" old = tmp_path / store.MODEL_NAME old.write_bytes(b"old invalid file") @@ -83,7 +83,26 @@ def test_interrupted_download_preserves_old_file(tmp_path, monkeypatch, 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")) + assert (tmp_path / f".{store.MODEL_NAME}.part").read_bytes() == artifact[:5] + + +def test_interrupted_download_resumes_with_http_range(tmp_path, monkeypatch, artifact): + """Resume a retained partial artifact instead of downloading it from zero.""" + partial = tmp_path / f".{store.MODEL_NAME}.part" + partial.write_bytes(artifact[:5]) + + class Resumed(io.BytesIO): + status = 206 + headers = {"Content-Range": f"bytes 5-{len(artifact) - 1}/{len(artifact)}"} + + def response(request, timeout): + assert request.headers["Range"] == "bytes=5-" + return Resumed(artifact[5:]) + + monkeypatch.setattr(store, "urlopen", response) + target = store.ensure_model(tmp_path) + assert target.read_bytes() == artifact + assert not partial.exists() def test_corrupt_cache_replaced_only_after_verification( diff --git a/tests/test_setup_shell.py b/tests/test_setup_shell.py index a22b923..7f2f255 100644 --- a/tests/test_setup_shell.py +++ b/tests/test_setup_shell.py @@ -86,12 +86,12 @@ def test_torch_backend_index_is_unambiguous(tmp_path, gpu, mirror): 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 + assert ["-m", "lama_inpaint", "download", "--verbose"] 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 ["-m", "lama_inpaint", "download", "--verbose"] not in commands assert "Setup complete!" not in result.stdout