feat: add progress bar for model downloads and update deployment files

download_models.py

@@ -1,10 +1,33 @@
 import os
 from azure.storage.blob import BlobServiceClient
 from dotenv import load_dotenv
+from tqdm import tqdm
 
 # Load environment variables
 load_dotenv()
 
+def download_with_progress(blob_client, local_path):
+    """Download blob with progress bar"""
+    # Get blob properties for total size
+    properties = blob_client.get_blob_properties()
+    total_size = properties.size
+
+    # Create progress bar
+    progress = tqdm(
+        total=total_size,
+        unit='iB',
+        unit_scale=True,
+        desc=f"Downloading {os.path.basename(local_path)}"
+    )
+
+    # Download in chunks
+    with open(local_path, "wb") as file:
+        download_stream = blob_client.download_blob()
+        for chunk in download_stream.chunks():
+            file.write(chunk)
+            progress.update(len(chunk))
+    progress.close()
+
 def download_from_azure():
     """Download models from Azure Blob Storage"""
     try:
@@ -15,6 +38,7 @@ def download_from_azure():
             
         container_name = os.getenv('AZURE_STORAGE_CONTAINER_NAME', 'pnid-models')
 
+        print("Connecting to Azure Blob Storage...")
         # Create the BlobServiceClient
         blob_service_client = BlobServiceClient.from_connection_string(connect_str)
         container_client = blob_service_client.get_container_client(container_name)
@@ -38,20 +62,21 @@ def download_from_azure():
             if not local_path:
                 continue
                 
-            print(f"Downloading {blob_name}...")
+            print(f"\nChecking {blob_name}...")
             os.makedirs(os.path.dirname(local_path), exist_ok=True)
             
             if not os.path.exists(local_path) or os.path.getsize(local_path) == 0:
                 blob_client = container_client.get_blob_client(blob_name)
-                with open(local_path, "wb") as file:
-                    data = blob_client.download_blob()
-                    file.write(data.readall())
-                print(f"Downloaded {blob_name} to {local_path}")
+                print(f"Starting download of {blob_name}...")
+                download_with_progress(blob_client, local_path)
+                print(f"Successfully downloaded {blob_name}")
             else:
                 print(f"Skipping {blob_name}, already exists")
 
+        print("\nAll models downloaded successfully!")
+
     except Exception as e:
-        print(f"Error downloading models from Azure: {str(e)}")
+        print(f"\nError downloading models from Azure: {str(e)}")
         raise
 
 if __name__ == "__main__":