Spaces:
Sleeping
Sleeping
| 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() |