KarelWintersky / download_hf_model.py

#!/usr/bin/env python3 """ Download all files from a HuggingFace model repository with retries, resume support, and integrity verification. Usage: python3 download_hf_model.py <repo_id> [repo_id ...] Examples: python3 download_hf_model.py ai-sage/GigaAM-Multilingual python3 download_hf_model.py meta-llama/Llama-3-8B-Instruct Each repo is downloaded into a directory named by replacing '/' with '.' in the repo_id (e.g. ai-sage.GigaAM-Multilingual). Adapted from download_lilith_q8.py. """ import argparse import hashlib import importlib import os import subprocess import sys import tempfile import time from pathlib import Path MAX_RETRIES = 10 RETRY_DELAY_BASE = 15 # seconds, doubles each retry BASE_DIR = Path(__file__).resolve().parent RUNTIME_DIR = BASE_DIR / ".hf_download_runtime" HF_HOME_DIR = RUNTIME_DIR / "hf-home" HF_HUB_CACHE_DIR = HF_HOME_DIR / "hub" HF_ASSETS_CACHE_DIR = HF_HOME_DIR / "assets" XDG_CACHE_HOME_DIR = RUNTIME_DIR / "xdg-cache" PIP_CACHE_DIR = RUNTIME_DIR / "pip-cache" PYTHON_VENDOR_DIR = RUNTIME_DIR / "python-vendor" TMP_DIR = RUNTIME_DIR / "tmp" def prepare_runtime_dirs() -> None: for path in ( RUNTIME_DIR, HF_HOME_DIR, HF_HUB_CACHE_DIR, HF_ASSETS_CACHE_DIR, XDG_CACHE_HOME_DIR, PIP_CACHE_DIR, PYTHON_VENDOR_DIR, TMP_DIR, ): path.mkdir(parents=True, exist_ok=True) os.environ["HF_HOME"] = str(HF_HOME_DIR) os.environ["HF_HUB_CACHE"] = str(HF_HUB_CACHE_DIR) os.environ["HF_ASSETS_CACHE"] = str(HF_ASSETS_CACHE_DIR) os.environ["XDG_CACHE_HOME"] = str(XDG_CACHE_HOME_DIR) os.environ["PIP_CACHE_DIR"] = str(PIP_CACHE_DIR) os.environ["HF_HUB_DISABLE_XET"] = "1" os.environ["HF_XET_RECONSTRUCT_WRITE_SEQUENTIALLY"] = "1" os.environ["TMPDIR"] = str(TMP_DIR) os.environ["TMP"] = str(TMP_DIR) os.environ["TEMP"] = str(TMP_DIR) tempfile.tempdir = str(TMP_DIR) if str(PYTHON_VENDOR_DIR) not in sys.path: sys.path.insert(0, str(PYTHON_VENDOR_DIR)) prepare_runtime_dirs() def parse_major_minor(version: str) -> tuple[int, int]: parts = [] for chunk in version.split("."): digits = "".join(ch for ch in chunk if ch.isdigit()) if digits: parts.append(int(digits)) if len(parts) == 2: break while len(parts) < 2: parts.append(0) return parts[0], parts[1] def install_huggingface_hub() -> None: print("Installing newer huggingface_hub into the local runtime dir...") subprocess.check_call( [ sys.executable, "-m", "pip", "install", "--upgrade", "--target", str(PYTHON_VENDOR_DIR), "--cache-dir", str(PIP_CACHE_DIR), "huggingface_hub>=1.0,<2", ] ) try: import huggingface_hub as _huggingface_hub if parse_major_minor(_huggingface_hub.__version__) < (1, 0): raise ImportError(f"Found old huggingface_hub {_huggingface_hub.__version__}") from huggingface_hub import HfApi, hf_hub_url from huggingface_hub.utils import RepositoryNotFoundError except Exception: install_huggingface_hub() importlib.invalidate_caches() for name in list(sys.modules): if name == "huggingface_hub" or name.startswith("huggingface_hub."): del sys.modules[name] from huggingface_hub import HfApi, hf_hub_url from huggingface_hub.utils import RepositoryNotFoundError def load_token() -> str | None: token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") if token: return token token_files = [ HF_HOME_DIR / "token", Path.home() / ".cache" / "huggingface" / "token", ] for token_file in token_files: if token_file.exists(): return token_file.read_text().strip() return None TOKEN = load_token() if not TOKEN: print("HF token not found. Proceeding without authentication (repo may be public).") else: print(f"Using HF token: {TOKEN[:8]}...") def repo_to_dirname(repo_id: str) -> str: return repo_id.replace("/", ".") def get_dir_size(path: Path) -> int: return sum(f.stat().st_size for f in path.rglob("*") if f.is_file()) def expected_size_for(sibling) -> int | None: size = getattr(sibling, "size", None) if size is not None: return size if sibling.lfs and sibling.lfs.get("size") is not None: return sibling.lfs["size"] return None def format_size(n: int | None) -> str: if n is None: return "unknown size" if n >= 1e9: return f"{n / 1e9:.2f} GB" if n >= 1e6: return f"{n / 1e6:.2f} MB" return f"{n / 1e3:.2f} KB" def quarantine_corrupt_files(paths: list[Path], download_dir: Path) -> None: if not paths: return quarantine_dir = download_dir / ".corrupt" quarantine_dir.mkdir(exist_ok=True) stamp = int(time.time()) for path in paths: target = quarantine_dir / f"{path.name}.{stamp}.corrupt" suffix = 1 while target.exists(): target = quarantine_dir / f"{path.name}.{stamp}.{suffix}.corrupt" suffix += 1 path.replace(target) print(f"Moved corrupt file aside: {target}") def download_file(rel_path: str, dest: Path, repo_id: str) -> None: url = hf_hub_url(repo_id=repo_id, filename=rel_path) cmd = [ "curl", "-L", "--fail", "--show-error", "--connect-timeout", "30", "--progress-bar", "-C", "-", "-o", str(dest), url, ] if TOKEN: cmd[1:1] = ["-H", f"Authorization: Bearer {TOKEN}"] subprocess.check_call(cmd) def download_repo_with_retries(repo_id: str, download_dir: Path) -> Path: """Download all files from a repo with curl, resume, and exponential backoff.""" for attempt in range(1, MAX_RETRIES + 1): try: print(f"\n{'=' * 60}") print(f"Download attempt {attempt}/{MAX_RETRIES}") print(f"{'=' * 60}") api = HfApi(token=TOKEN) repo_info = api.repo_info(repo_id=repo_id, token=TOKEN, files_metadata=True) siblings = [s for s in (repo_info.siblings or []) if s.rfilename] if attempt == 1: # Disk space check needed = 0 for s in siblings: exp = expected_size_for(s) dest = download_dir / s.rfilename have = dest.stat().st_size if dest.exists() else 0 if exp is not None: needed += max(exp - have, 0) try: stat = os.statvfs(download_dir) free = stat.f_bavail * stat.f_frsize print(f"Disk check: need ~{format_size(needed)} more, free ~{format_size(free)} on destination.") if needed > free: print("WARNING: not enough free space for the remaining download.") except OSError: pass siblings = sorted( siblings, key=lambda s: expected_size_for(s) or 0, reverse=True, ) size_before = get_dir_size(download_dir) if download_dir.exists() else 0 t0 = time.time() for index, sibling in enumerate(siblings, start=1): rel_path = sibling.rfilename dest = download_dir / rel_path dest.parent.mkdir(parents=True, exist_ok=True) expected_size = expected_size_for(sibling) current_size = dest.stat().st_size if dest.exists() else 0 if expected_size is not None and dest.exists() and current_size == expected_size: print( f"[{index}/{len(siblings)}] Already present: {rel_path} " f"({format_size(expected_size)})" ) continue if expected_size is not None and dest.exists() and current_size > expected_size: quarantine_corrupt_files([dest], download_dir) current_size = 0 if current_size > 0: print( f"[{index}/{len(siblings)}] Resuming: {rel_path} " f"from {format_size(current_size)} / {format_size(expected_size)}" ) else: print(f"[{index}/{len(siblings)}] Downloading: {rel_path} ({format_size(expected_size)})") download_file(rel_path, dest, repo_id) if expected_size is not None: actual_size = dest.stat().st_size if actual_size != expected_size: raise ValueError( f"Size mismatch after download for {rel_path}: " f"expected {expected_size}, got {actual_size}" ) elapsed = time.time() - t0 size_after = get_dir_size(download_dir) downloaded = max(size_after - size_before, 0) speed = downloaded / elapsed if elapsed > 0 else 0 print(f"\nDownload complete: {download_dir}") print(f" Downloaded: {format_size(downloaded)} in {elapsed:.0f}s") print(f" Avg speed: {speed / 1e6:.1f} MB/s") return download_dir except KeyboardInterrupt: print("\nInterrupted by user.") sys.exit(1) except (RepositoryNotFoundError, FileNotFoundError) as exc: print(f"\nFatal error: {exc}") sys.exit(1) except Exception as exc: delay = RETRY_DELAY_BASE * (2 ** (attempt - 1)) print(f"\nError on attempt {attempt}: {exc}") if attempt < MAX_RETRIES: print(f"Retrying in {delay}s...") time.sleep(delay) else: print("All retries exhausted.") sys.exit(1) def verify_integrity(repo_id: str, download_dir: Path) -> tuple[bool, list[Path]]: """Verify all downloaded files against repo metadata.""" print(f"\n{'=' * 60}") print("Verifying file integrity...") print(f"{'=' * 60}") api = HfApi(token=TOKEN) try: repo_info = api.repo_info(repo_id=repo_id, token=TOKEN, files_metadata=True) except Exception as exc: print(f"Could not fetch repo metadata for verification: {exc}") return False, [] siblings = [s for s in (repo_info.siblings or []) if s.rfilename] errors = [] checked = 0 corrupt_files: list[Path] = [] for sibling in siblings: if not sibling.rfilename: continue file_path = download_dir / sibling.rfilename if not file_path.exists(): errors.append(f"MISSING: {sibling.rfilename}") continue expected_sha = None hash_type = None if sibling.lfs and sibling.lfs.get("sha256"): expected_sha = sibling.lfs["sha256"] hash_type = "sha256" elif hasattr(sibling, "blob_id") and sibling.blob_id: expected_sha = sibling.blob_id hash_type = "git-sha1" else: continue checked += 1 size = file_path.stat().st_size print(f" Checking {sibling.rfilename} ({format_size(size)})...", end=" ", flush=True) if hash_type == "sha256": digest = hashlib.sha256() with open(file_path, "rb") as handle: while chunk := handle.read(8 * 1024 * 1024): digest.update(chunk) actual = digest.hexdigest() else: digest = hashlib.sha1() digest.update(f"blob {size}\0".encode()) with open(file_path, "rb") as handle: while chunk := handle.read(8 * 1024 * 1024): digest.update(chunk) actual = digest.hexdigest() if actual == expected_sha: print("OK") else: print("MISMATCH") corrupt_files.append(file_path) errors.append( f"HASH MISMATCH: {sibling.rfilename}\n" f" expected: {expected_sha}\n" f" actual: {actual}" ) print(f"\nChecked {checked} file(s).") if errors: print("\nIntegrity errors:") for err in errors: print(f" - {err}") return False, corrupt_files print("All files verified successfully.") return True, [] def main() -> None: parser = argparse.ArgumentParser( description="Download files from HuggingFace model repositories with retry and verification." ) parser.add_argument( "repos", nargs="+", help="HuggingFace repo IDs to download (e.g. ai-sage/GigaAM-Multilingual)", ) args = parser.parse_args() for repo_id in args.repos: dirname = repo_to_dirname(repo_id) download_dir = BASE_DIR / dirname download_dir.mkdir(parents=True, exist_ok=True) print(f"\n{'#' * 60}") print(f"# Repo: {repo_id}") print(f"# Target: {download_dir}") print(f"{'#' * 60}") download_repo_with_retries(repo_id, download_dir) ok, corrupt_files = verify_integrity(repo_id, download_dir) if not ok: quarantine_corrupt_files(corrupt_files, download_dir) print(f"\nIntegrity check FAILED for {repo_id}. Re-run the script to resume.") sys.exit(1) print(f"\nDone. Model saved to: {download_dir}") if __name__ == "__main__": main()