app.py

import os
import requests
import gradio as gr
from huggingface_hub import HfApi, login, create_repo
import tempfile

def download_and_push_model(progress=gr.Progress()):
    """
    Download SAM model and push it to Hugging Face Space
    """
    try:
        # Login to Hugging Face
        token = os.environ.get('HF_TOKEN')
        if not token:
            return "❌ Error: HF_TOKEN not found in environment variables. Please add it to Space secrets."
        
        login(token)  # Authenticate with Hugging Face
        
        # Initialize Hugging Face API
        api = HfApi()
        space_id = "lyimo/downloadmodel"
        model_url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
        
        progress(0.05, desc="Ensuring repository exists...")
        try:
            # Try to create the repo (will fail if it already exists, which is fine)
            create_repo(
                repo_id=space_id,
                repo_type="space",
                token=token,
                exist_ok=True
            )
        except Exception as e:
            progress(0.1, desc="Repository already exists, continuing...")
            pass
        
        # Create a temporary directory for downloading
        with tempfile.TemporaryDirectory() as temp_dir:
            progress(0.1, desc="Started download process...")
            
            # Download file
            local_path = os.path.join(temp_dir, "sam_vit_h_4b8939.pth")
            response = requests.get(model_url, stream=True)
            total_size = int(response.headers.get('content-length', 0))
            
            progress(0.2, desc="Downloading model...")
            downloaded_size = 0
            
            with open(local_path, 'wb') as file:
                for chunk in response.iter_content(chunk_size=1024*1024):
                    if chunk:
                        file.write(chunk)
                        downloaded_size += len(chunk)
                        progress(0.2 + 0.6 * (downloaded_size/total_size), 
                               desc=f"Downloading... {downloaded_size/(1024*1024):.1f}MB / {total_size/(1024*1024):.1f}MB")
            
            progress(0.8, desc="Uploading to Hugging Face Space...")
            
            # Upload to Hugging Face using commit operation
            api.upload_file(
                path_or_fileobj=local_path,
                path_in_repo="sam_vit_h_4b8939.pth",
                repo_id=space_id,
                repo_type="space",
                token=token,
                commit_message="Upload SAM model weights"
            )
            
            progress(1.0, desc="Complete!")
            return "✅ Model successfully downloaded and pushed to your Space!"
            
    except Exception as e:
        return f"❌ Error: {str(e)}\nToken status: {'Token exists' if token else 'No token found'}"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Download SAM Model to Space")
    gr.Markdown("This will download the Segment Anything Model (SAM) weights (~2.4GB) and push to this Space")
    
    with gr.Row():
        download_button = gr.Button("📥 Download & Push Model", variant="primary")
        status_text = gr.Textbox(label="Status", interactive=False)
    
    gr.Markdown("""
    ### Important Setup Steps:
    1. Get your Hugging Face token from https://huggingface.co/settings/tokens
    2. Add the token to Space secrets:
       - Go to Space Settings > Secrets
       - Add new secret named `HF_TOKEN`
       - Paste your token as the value
    3. Restart the Space after adding the token
    """)
    
    download_button.click(
        fn=download_and_push_model,
        outputs=status_text,
        show_progress=True,
    )
    
    gr.Markdown("""
    ### Notes:
    - Download size is approximately 2.4GB
    - The model will be pushed to the Space repository
    - Please wait for both download and upload to complete
    - You can check the files tab after completion
    """)

if __name__ == "__main__":
    demo.launch()