File size: 3,724 Bytes
f3972ea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
#!/usr/bin/env python3
"""
Script to upload all .pt and .onnx model files from models_mask directory to Hugging Face.
"""
import os
import glob
from pathlib import Path
from huggingface_hub import HfApi, login
from tqdm import tqdm
def upload_models_mask():
"""
Upload all .pt and .onnx files from models_mask directory to Hugging Face
"""
# Configuration
repo_name = "PetchMa/BLADE_FRBNN" # Change this to your desired repo name
models_dir = "models_mask"
# Initialize HF API
api = HfApi()
# Login to Hugging Face (you need to have your token set up)
print("Logging into Hugging Face...")
try:
login()
print("✓ Successfully logged into Hugging Face")
except Exception as e:
print(f"❌ Failed to login to Hugging Face: {e}")
print("Please make sure you have your HF token set up:")
print("1. Run: huggingface-cli login")
print("2. Or set HF_TOKEN environment variable")
return
# Check if models directory exists
if not os.path.exists(models_dir):
print(f"❌ Directory {models_dir} does not exist!")
return
# Find all .pt and .onnx files
pt_files = glob.glob(os.path.join(models_dir, "*.pt"))
onnx_files = glob.glob(os.path.join(models_dir, "*.onnx"))
all_files = pt_files + onnx_files
if not all_files:
print(f"❌ No .pt or .onnx files found in {models_dir}")
return
print(f"Found {len(all_files)} model files to upload:")
print(f" - {len(pt_files)} .pt files")
print(f" - {len(onnx_files)} .onnx files")
# Calculate total size
total_size = sum(os.path.getsize(f) for f in all_files)
total_size_gb = total_size / (1024**3)
print(f" - Total size: {total_size_gb:.2f} GB")
# Create repository if it doesn't exist
try:
api.create_repo(repo_name, exist_ok=True)
print(f"✓ Repository {repo_name} is ready")
except Exception as e:
print(f"❌ Failed to create repository: {e}")
return
# Use upload_large_folder for better performance with many files
print(f"\nUploading {models_dir} directory to {repo_name}...")
print("This may take a while due to the large number of files...")
try:
api.upload_large_folder(
folder_path=models_dir,
repo_id=repo_name,
repo_type="model",
commit_message=f"Upload models_mask directory with {len(all_files)} model files",
allow_patterns=["*.pt", "*.onnx"], # Only upload .pt and .onnx files
ignore_patterns=["*.png", "*.pdf", "*.pkl", "*.bin"], # Ignore other file types
)
print(f"\n✓ Successfully uploaded {len(all_files)} model files!")
print(f"✓ Repository: https://huggingface.co/{repo_name}")
print(f"✓ Models available at: https://huggingface.co/{repo_name}/tree/main/{models_dir}")
except Exception as e:
print(f"❌ Failed to upload folder: {e}")
print("You might want to try uploading in smaller batches or check your internet connection.")
def main():
"""
Main function to run the upload process
"""
print("BLADE FRBNN Model Upload Tool")
print("=" * 40)
print("This script will upload all .pt and .onnx files from models_mask/")
print("to your Hugging Face repository using the efficient large folder upload.")
print()
# Ask for confirmation
response = input("Do you want to proceed? (y/N): ")
if response.lower() not in ['y', 'yes']:
print("Upload cancelled.")
return
upload_models_mask()
if __name__ == "__main__":
main() |