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