File size: 4,068 Bytes
8dfc891 aff0043 7973235 aff0043 b86605e aff0043 b86605e aff0043 7cff9dc aff0043 b86605e aff0043 b86605e aff0043 a07334a aff0043 a07334a aff0043 7c43484 aff0043 5944a16 aff0043 7973235 5c02e21 a07334a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
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/tobacco"
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() |