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()