Update app.py

app.py

@@ -1,59 +1,76 @@
 import os
 import requests
-from tqdm import tqdm
-import logging
 import gradio as gr
+import time
 
-def download_sam_model():
+def download_sam_model(progress=gr.Progress()):
     """
-    Download the SAM model weights to the Hugging Face Space repository root
+    Download the SAM model weights with visible progress indicator
     """
     url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
     
     try:
-        # Get the repository root directory (same level as app.py)
+        # 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)
         
-        # Setup logging
-        logging.basicConfig(level=logging.INFO)
-        logger = logging.getLogger(__name__)
-        
         # Check if file already exists
         if os.path.exists(save_path):
-            return gr.Info("Model file already exists in the repository root")
+            return "Model file already exists in the repository root"
             
-        # Send a HEAD request to get file size
+        # Get file size
         response = requests.head(url)
-        file_size = int(response.headers.get('content-length', 0))
+        total_size = int(response.headers.get('content-length', 0))
         
-        # Download the file with progress bar
-        logger.info(f"Downloading {filename} to Space repository root")
+        # Download with progress tracking
         response = requests.get(url, stream=True)
-        progress = tqdm(total=file_size, unit='iB', unit_scale=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=1024):
-                progress.update(len(data))
+            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")
         
-        progress.close()
-        
-        if file_size != 0 and progress.n != file_size:
-            raise Exception("Downloaded file size does not match expected size")
+        if downloaded_size != total_size:
+            if os.path.exists(save_path):
+                os.remove(save_path)
+            return "Error: Download incomplete. Please try again."
             
-        return gr.Info("Model downloaded successfully!")
+        return f"✅ Model successfully downloaded! Saved as {filename}"
         
     except Exception as e:
-        logger.error(f"Error downloading file: {str(e)}")
-        return gr.Error(f"Failed to download model: {str(e)}")
+        if os.path.exists(save_path):
+            os.remove(save_path)
+        return f"❌ Error downloading file: {str(e)}"
 
-# Create a simple Gradio interface
+# Create a Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Download SAM Model")
-    download_button = gr.Button("Download Model")
-    download_button.click(fn=download_sam_model)
+    gr.Markdown("This will download the Segment Anything Model (SAM) weights (~2.4GB)")
+    
+    with gr.Row():
+        download_button = gr.Button("📥 Download Model", variant="primary")
+        status_text = gr.Textbox(label="Status", interactive=False)
+    
+    download_button.click(
+        fn=download_sam_model,
+        outputs=status_text,
+        show_progress=True,
+    )
+    
+    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
+    """)
 
 if __name__ == "__main__":
     demo.launch()