Add Hugging Face Spaces upload functionality

- Add direct upload to HF Spaces
- Improve download progress tracking
- Add better error handling
- Implement smaller chunk sizes for downloads

download_models.py

@@ -2,51 +2,80 @@ import os
 from azure.storage.blob import BlobServiceClient
 from dotenv import load_dotenv
 from tqdm import tqdm
+import requests
+from huggingface_hub import HfApi, upload_file
 
 # Load environment variables
 load_dotenv()
 
+def upload_to_hf_spaces(local_path, repo_id, token):
+    """Upload file to Hugging Face Spaces"""
+    try:
+        api = HfApi()
+        print(f"\nUploading {os.path.basename(local_path)} to {repo_id}...")
+        response = api.upload_file(
+            path_or_fileobj=local_path,
+            path_in_repo=os.path.basename(local_path),
+            repo_id=repo_id,
+            token=token
+        )
+        print(f"Successfully uploaded to {response}")
+        return True
+    except Exception as e:
+        print(f"Error uploading to Hugging Face: {str(e)}")
+        return False
+
 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"""
+    """Download blob with improved progress tracking"""
+    try:
+        properties = blob_client.get_blob_properties()
+        total_size = properties.size
+        
+        # Configure progress bar
+        progress = tqdm(
+            total=total_size,
+            unit='B',
+            unit_scale=True,
+            desc=f"Downloading {os.path.basename(local_path)}",
+            ncols=80
+        )
+
+        # Download in smaller chunks (5MB)
+        chunk_size = 5 * 1024 * 1024  # 5MB chunks
+        with open(local_path, "wb") as file:
+            download_stream = blob_client.download_blob()
+            for chunk in download_stream.chunks(chunk_size=chunk_size):
+                if chunk:
+                    file.write(chunk)
+                    progress.update(len(chunk))
+        
+        progress.close()
+        return True
+    except Exception as e:
+        print(f"\nError during download: {str(e)}")
+        # Clean up partial download
+        if os.path.exists(local_path):
+            os.remove(local_path)
+        return False
+
+def download_from_azure(upload_to_hf=False):
+    """Download models from Azure Blob Storage with optional HF upload"""
     try:
-        # Try primary connection first, fallback to secondary
+        # Get connection strings
         connect_str = os.getenv('AZURE_STORAGE_PRIMARY_CONNECTION')
         if not connect_str:
             connect_str = os.getenv('AZURE_STORAGE_SECONDARY_CONNECTION')
-            
+        
         container_name = os.getenv('AZURE_STORAGE_CONTAINER_NAME', 'pnid-models')
+        hf_token = os.getenv('HF_TOKEN')
+        hf_repo = os.getenv('HF_REPO_ID')
 
         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)
 
-        # Create models directory
         os.makedirs('models', exist_ok=True)
 
-        # Define model files to download
         model_files = {
             'Intui_SDM_41.pt': os.getenv('MODEL_SDM_41_PATH'),
             'Intui_SDM_30.pt': os.getenv('MODEL_SDM_30_PATH'),
@@ -57,27 +86,41 @@ def download_from_azure():
             'intui_LDM_01.pt': os.getenv('MODEL_LDM_PATH')
         }
 
-        # Download each model
         for blob_name, local_path in model_files.items():
             if not local_path:
                 continue
-                
-            print(f"\nChecking {blob_name}...")
+
+            print(f"\nProcessing {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:
+
+            # Check if file needs to be downloaded
+            needs_download = (
+                not os.path.exists(local_path) or 
+                os.path.getsize(local_path) == 0
+            )
+
+            if needs_download:
                 blob_client = container_client.get_blob_client(blob_name)
-                print(f"Starting download of {blob_name}...")
-                download_with_progress(blob_client, local_path)
+                success = download_with_progress(blob_client, local_path)
+                
+                if not success:
+                    print(f"Failed to download {blob_name}")
+                    continue
+                
                 print(f"Successfully downloaded {blob_name}")
             else:
                 print(f"Skipping {blob_name}, already exists")
 
-        print("\nAll models downloaded successfully!")
+            # Upload to HF Spaces if requested
+            if upload_to_hf and hf_token and hf_repo:
+                upload_to_hf_spaces(local_path, hf_repo, hf_token)
+
+        print("\nAll operations completed!")
 
     except Exception as e:
-        print(f"\nError downloading models from Azure: {str(e)}")
+        print(f"\nError in main process: {str(e)}")
         raise
 
 if __name__ == "__main__":
-    download_from_azure()
+    # Set to True if you want to upload to HF Spaces
+    download_from_azure(upload_to_hf=False)