download_weights.py

#!/usr/bin/env python3
"""
Download FoundationPose pre-trained model weights from Hugging Face.

Weights can be hosted in a HF model repository (recommended) or downloaded
manually from the official Google Drive.
"""

import os
import sys
from pathlib import Path

try:
    from huggingface_hub import hf_hub_download, snapshot_download
except ImportError:
    print("Installing huggingface_hub...")
    os.system(f"{sys.executable} -m pip install huggingface_hub")
    from huggingface_hub import hf_hub_download, snapshot_download


# Configuration
HF_MODEL_REPO = os.environ.get("FOUNDATIONPOSE_MODEL_REPO", "gpue/foundationpose-weights")
USE_HF_WEIGHTS = os.environ.get("USE_HF_WEIGHTS", "true").lower() == "true"


def download_from_huggingface(weights_dir: Path) -> bool:
    """Download weights from Hugging Face model repository.

    Args:
        weights_dir: Directory to save weights

    Returns:
        True if successful
    """
    print("=" * 60)
    print("Downloading from Hugging Face Model Repository")
    print("=" * 60)
    print(f"Repository: {HF_MODEL_REPO}")
    print(f"Target: {weights_dir.absolute()}")
    print()

    try:
        # Get HF token if available (for private repos)
        hf_token = os.environ.get("HF_TOKEN")
        if hf_token:
            print("🔒 Using HF_TOKEN for authentication (private repository)")

        print("Downloading model weights...")
        print("(This may take several minutes on first run)")
        print()

        # Download entire repository
        snapshot_download(
            repo_id=HF_MODEL_REPO,
            local_dir=str(weights_dir),
            local_dir_use_symlinks=False,
            resume_download=True,
            token=hf_token  # Will use token if provided, None otherwise
        )

        print()
        print("✓ Download complete!")
        return True

    except Exception as e:
        error_msg = str(e).lower()
        print(f"✗ Download failed: {e}")
        print()

        # Check if it's an authentication error
        if "401" in error_msg or "403" in error_msg or "authentication" in error_msg or "token" in error_msg:
            print("🔒 Authentication Error - Repository is private!")
            print()
            print("Solutions:")
            print("  Option 1: Make repository public")
            print(f"    Visit: https://huggingface.co/{HF_MODEL_REPO}/settings")
            print("    Change visibility to 'Public'")
            print()
            print("  Option 2: Add HF token to Space secrets")
            print("    1. Get token: https://huggingface.co/settings/tokens")
            print("    2. Add to Space secrets as 'HF_TOKEN'")
            print()
        else:
            print("Possible issues:")
            print(f"  1. Repository '{HF_MODEL_REPO}' doesn't exist")
            print("  2. Repository is private (need HF_TOKEN in secrets)")
            print("  3. Network error")
            print()
            print("To create the model repository:")
            print("  1. Visit: https://huggingface.co/new")
            print("  2. Create a model repo (e.g., 'gpue/foundationpose-weights')")
            print("  3. Upload weights using:")
            print("     huggingface-cli upload gpue/foundationpose-weights ./weights/")
            print()
        return False


def manual_download_instructions(weights_dir: Path):
    """Print instructions for manual weight download."""
    print("=" * 60)
    print("Manual Weight Download Instructions")
    print("=" * 60)
    print()
    print("Option 1: Download from official Google Drive")
    print("-" * 40)
    print("1. Visit: https://drive.google.com/drive/folders/1GCyGE-LbFGgRC-FuGsF3a1zeBuzsQ1Da")
    print("2. Download these folders:")
    print("   - 2023-10-28-18-33-37/ (refiner weights)")
    print("   - 2024-01-11-20-02-45/ (scorer weights)")
    print(f"3. Extract to: {weights_dir.absolute()}")
    print()
    print("Option 2: Create Hugging Face model repository")
    print("-" * 40)
    print("1. Download weights from Google Drive (see above)")
    print("2. Create HF model repo: https://huggingface.co/new")
    print("3. Upload weights:")
    print("   pip install huggingface_hub")
    print("   huggingface-cli login")
    print(f"   huggingface-cli upload YOUR_USERNAME/foundationpose-weights {weights_dir}/")
    print("4. Set environment variable:")
    print(f"   export FOUNDATIONPOSE_MODEL_REPO=YOUR_USERNAME/foundationpose-weights")
    print()


def check_weights_exist(weights_dir: Path) -> bool:
    """Check if weights already exist locally.

    Args:
        weights_dir: Directory containing weights

    Returns:
        True if weights exist
    """
    required_folders = [
        weights_dir / "2023-10-28-18-33-37",
        weights_dir / "2024-01-11-20-02-45"
    ]

    return all(folder.exists() and any(folder.iterdir()) for folder in required_folders)


def download_weights() -> bool:
    """Download or check for FoundationPose weights.

    Returns:
        True if weights are available
    """
    weights_dir = Path("weights")
    weights_dir.mkdir(exist_ok=True)

    # Check if weights already exist
    if check_weights_exist(weights_dir):
        print("✓ Model weights found locally!")
        print(f"  Location: {weights_dir.absolute()}")
        return True

    print("Model weights not found locally.")
    print()

    # Try downloading from Hugging Face
    if USE_HF_WEIGHTS:
        print(f"Attempting to download from Hugging Face...")
        print(f"Repository: {HF_MODEL_REPO}")
        print()

        if download_from_huggingface(weights_dir):
            return True

        print()
        print("Hugging Face download failed. See manual instructions below.")
        print()

    # Show manual instructions
    manual_download_instructions(weights_dir)

    return False


if __name__ == "__main__":
    print()
    success = download_weights()
    print()

    if success:
        print("=" * 60)
        print("✓ Ready to use FoundationPose!")
        print("=" * 60)
        sys.exit(0)
    else:
        print("=" * 60)
        print("⚠ Weights not available")
        print("=" * 60)
        print()
        print("Space will run in PLACEHOLDER mode.")
        print("To enable real inference, follow instructions above.")
        sys.exit(1)