verify_weights.py

#!/usr/bin/env python3
"""
Verify that FoundationPose weights are accessible from Hugging Face model repository.

This script checks:
1. Model repo exists and is accessible
2. Required weight files are present
3. Files can be downloaded successfully
"""

import os
import sys
from pathlib import Path

try:
    from huggingface_hub import list_repo_files, hf_hub_download
except ImportError:
    print("❌ huggingface_hub not installed")
    print("Install with: pip install huggingface_hub")
    sys.exit(1)

# Configuration
MODEL_REPO = os.environ.get("FOUNDATIONPOSE_MODEL_REPO", "gpue/foundationpose-weights")

# Required files
REQUIRED_FILES = [
    "2023-10-28-18-33-37/config.yml",
    "2023-10-28-18-33-37/model_best.pth",
    "2024-01-11-20-02-45/config.yml",
    "2024-01-11-20-02-45/model_best.pth",
]


def verify_repo_access():
    """Verify model repository is accessible."""
    print(f"Checking repository: {MODEL_REPO}")
    print("-" * 60)

    try:
        files = list_repo_files(repo_id=MODEL_REPO, repo_type="model")
        print(f"✓ Repository accessible")
        print(f"✓ Found {len(files)} files")
        return files
    except Exception as e:
        print(f"❌ Cannot access repository: {e}")
        return None


def verify_required_files(repo_files):
    """Verify all required weight files are present."""
    print("\nChecking required files:")
    print("-" * 60)

    all_present = True
    for required_file in REQUIRED_FILES:
        if required_file in repo_files:
            print(f"✓ {required_file}")
        else:
            print(f"❌ Missing: {required_file}")
            all_present = False

    return all_present


def test_download():
    """Test downloading a small file."""
    print("\nTesting download:")
    print("-" * 60)

    try:
        # Download a small config file to test connectivity
        test_file = "2023-10-28-18-33-37/config.yml"
        print(f"Downloading {test_file}...")

        downloaded = hf_hub_download(
            repo_id=MODEL_REPO,
            filename=test_file,
            repo_type="model"
        )

        print(f"✓ Download successful: {downloaded}")

        # Check file size
        size = Path(downloaded).stat().st_size
        print(f"✓ File size: {size:,} bytes")

        return True

    except Exception as e:
        print(f"❌ Download failed: {e}")
        return False


def main():
    """Run all verification checks."""
    print("=" * 60)
    print("FoundationPose Model Repository Verification")
    print("=" * 60)
    print()

    # Check 1: Repository access
    repo_files = verify_repo_access()
    if repo_files is None:
        print("\n❌ Verification failed: Cannot access repository")
        sys.exit(1)

    # Check 2: Required files
    has_all_files = verify_required_files(repo_files)
    if not has_all_files:
        print("\n❌ Verification failed: Missing required files")
        sys.exit(1)

    # Check 3: Download test
    can_download = test_download()
    if not can_download:
        print("\n❌ Verification failed: Cannot download files")
        sys.exit(1)

    # All checks passed
    print()
    print("=" * 60)
    print("✓ All verification checks passed!")
    print("=" * 60)
    print()
    print(f"Model repository '{MODEL_REPO}' is ready to use.")
    print()
    print("To use in your Space:")
    print("  1. Set environment variable:")
    print(f"     FOUNDATIONPOSE_MODEL_REPO={MODEL_REPO}")
    print("  2. Set USE_HF_WEIGHTS=true")
    print("  3. Set USE_REAL_MODEL=true")
    print()

    return 0


if __name__ == "__main__":
    sys.exit(main())