Update app.py

app.py

@@ -1,66 +1,66 @@
 import os
 import requests
 import gradio as gr
-import time
+from huggingface_hub import HfApi
+import tempfile
 
-def download_sam_model(progress=gr.Progress()):
+def download_and_push_model(progress=gr.Progress()):
     """
-    Download the SAM model weights with visible progress indicator
+    Download SAM model and push it to Hugging Face Space
     """
-    url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
-    
     try:
-        # Get the repository root directory
-        root_dir = os.path.dirname(os.path.abspath(__file__))
-        filename = url.split('/')[-1]
-        save_path = os.path.join(root_dir, filename)
+        # Initialize Hugging Face API
+        api = HfApi()
+        space_id = "lyimo/downloadmodel"
+        model_url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
         
-        # Check if file already exists
-        if os.path.exists(save_path):
-            return "Model file already exists in the repository root"
+        # Create a temporary directory for downloading
+        with tempfile.TemporaryDirectory() as temp_dir:
+            progress(0.1, desc="Started download process...")
             
-        # Get file size
-        response = requests.head(url)
-        total_size = int(response.headers.get('content-length', 0))
-        
-        # Download with progress tracking
-        response = requests.get(url, stream=True)
-        progress(0, desc="Starting download...")
-        
-        downloaded_size = 0
-        chunk_size = 1024 * 1024  # 1MB chunks
-        
-        with open(save_path, 'wb') as file:
-            for data in response.iter_content(chunk_size=chunk_size):
-                file.write(data)
-                downloaded_size += len(data)
-                # Update progress bar
-                progress(downloaded_size/total_size, 
-                        desc=f"Downloading... {downloaded_size/(1024*1024):.1f}MB / {total_size/(1024*1024):.1f}MB")
-        
-        if downloaded_size != total_size:
-            if os.path.exists(save_path):
-                os.remove(save_path)
-            return "Error: Download incomplete. Please try again."
+            # 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="Upload to Hugging Face Space...")
+            
+            # Upload to Hugging Face
+            api.upload_file(
+                path_or_fileobj=local_path,
+                path_in_repo="sam_vit_h_4b8939.pth",
+                repo_id=space_id,
+                repo_type="space"
+            )
+            
+            progress(1.0, desc="Complete!")
+            return "✅ Model successfully downloaded and pushed to your Space!"
             
-        return f"✅ Model successfully downloaded! Saved as {filename}"
-        
     except Exception as e:
-        if os.path.exists(save_path):
-            os.remove(save_path)
-        return f"❌ Error downloading file: {str(e)}"
+        return f"❌ Error: {str(e)}"
 
-# Create a Gradio interface
+# Create Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Download SAM Model")
-    gr.Markdown("This will download the Segment Anything Model (SAM) weights (~2.4GB)")
+    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 Model", variant="primary")
+        download_button = gr.Button("📥 Download & Push Model", variant="primary")
         status_text = gr.Textbox(label="Status", interactive=False)
     
     download_button.click(
-        fn=download_sam_model,
+        fn=download_and_push_model,
         outputs=status_text,
         show_progress=True,
     )
@@ -68,8 +68,9 @@ with gr.Blocks() as demo:
     gr.Markdown("""
     ### Notes:
     - Download size is approximately 2.4GB
-    - The model will be saved in the Space's root directory
-    - Please wait for the download to complete
+    - 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__":