#!/usr/bin/env python3
"""
Simple upload script for ISNet model to Hugging Face Hub
"""
import os
from huggingface_hub import HfApi, login

def upload_to_huggingface(repo_name, upload_dir):
    """Upload the model to Hugging Face Hub"""
    
    # Login to Hugging Face
    print("Please enter your Hugging Face token:")
    token = input().strip()
    login(token)
    
    api = HfApi()
    
    # Create the repository
    print(f"Creating repository: {repo_name}")
    api.create_repo(repo_name, exist_ok=True)
    
    # Upload all files
    print("Uploading model files...")
    for root, dirs, files in os.walk(upload_dir):
        for file in files:
            file_path = os.path.join(root, file)
            relative_path = os.path.relpath(file_path, upload_dir)
            
            with open(file_path, "rb") as f:
                api.upload_file(
                    path_or_fileobj=f,
                    path_in_repo=relative_path,
                    repo_id=repo_name
                )
            print(f"✅ Uploaded: {relative_path}")
    
    print(f"🎉 Model successfully uploaded to: https://huggingface.co/{repo_name}")

if __name__ == "__main__":
    print("🚀 Uploading ISNet model to Hugging Face...")
    
    # Check if model files exist
    upload_dir = "model_for_upload"
    if not os.path.exists(upload_dir):
        print(f"❌ Upload directory {upload_dir} not found!")
        exit(1)
    
    # List files to be uploaded
    print(f"📁 Files to upload from {upload_dir}:")
    for root, dirs, files in os.walk(upload_dir):
        for file in files:
            file_path = os.path.join(root, file)
            relative_path = os.path.relpath(file_path, upload_dir)
            print(f"  - {relative_path}")
    
    # Get repository name
    print("\nEnter your Hugging Face username:")
    username = input().strip()
    repo_name = f"{username}/isnet-background-remover"
    
    # Upload to Hugging Face
    upload_to_huggingface(repo_name, upload_dir)
    
    print(f"\n🎯 After upload, users can download your model with:")
    print(f"from transformers import AutoModelForImageSegmentation")
    print(f"model = AutoModelForImageSegmentation.from_pretrained('{repo_name}', trust_remote_code=True)")