downloadmodel / app.py
lyimo's picture
Update app.py
cd47295 verified
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()