Add verbose resumable LaMA downloads

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

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)

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

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)..."

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(

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