adding redirect for all requests

3ac868729e4972450a3c232a71fd588a36e82a64 · 2024-11-18 14:22:28 -0800 · cadenmackenzie
Files touched

M exo/download/hf/hf_helpers.py
Diff

commit 3ac868729e4972450a3c232a71fd588a36e82a64
Author: cadenmackenzie <cadenmackenzie@gmail.com>
Date:   Mon Nov 18 14:22:28 2024 -0800

    adding redirect for all requests
---
 exo/download/hf/hf_helpers.py | 43 +++++++++++++++++--------------------------
 1 file changed, 17 insertions(+), 26 deletions(-)

diff --git a/exo/download/hf/hf_helpers.py b/exo/download/hf/hf_helpers.py
index 01a2329b..c8e02a8d 100644
--- a/exo/download/hf/hf_helpers.py
+++ b/exo/download/hf/hf_helpers.py
@@ -448,33 +448,24 @@ async def get_file_download_percentage(
         url = urljoin(base_url, file_path)
         headers = await get_auth_headers()
         
-        # For safetensors files, we need to follow redirects and use GET instead of HEAD
-        if file_path.endswith('.safetensors'):
-            async with session.get(url, headers=headers, allow_redirects=True) as response:
-                if response.status != 200:
-                    if DEBUG >= 2: print(f"Failed to get remote file info for {file_path}: {response.status}")
-                    return 0
-                    
-                remote_size = int(response.headers.get('Content-Length', 0))
-                # Don't download the actual file, just get the headers
-                await response.release()
-        else:
-            async with session.head(url, headers=headers) as response:
-                if response.status != 200:
-                    if DEBUG >= 2: print(f"Failed to get remote file info for {file_path}: {response.status}")
-                    return 0
-                remote_size = int(response.headers.get('Content-Length', 0))
-        
-        if remote_size == 0:
-            if DEBUG >= 2: print(f"Remote size is 0 for {file_path}")
-            return 0
-            
-        # Only return 100% if sizes match exactly
-        if local_size == remote_size:
-            return 100.0
+        # Use HEAD request with redirect following for all files
+        async with session.head(url, headers=headers, allow_redirects=True) as response:
+            if response.status != 200:
+                if DEBUG >= 2: print(f"Failed to get remote file info for {file_path}: {response.status}")
+                return 0
+                
+            remote_size = int(response.headers.get('Content-Length', 0))
             
-        # Calculate percentage based on sizes
-        return (local_size / remote_size) * 100 if remote_size > 0 else 0
+            if remote_size == 0:
+                if DEBUG >= 2: print(f"Remote size is 0 for {file_path}")
+                return 0
+                
+            # Only return 100% if sizes match exactly
+            if local_size == remote_size:
+                return 100.0
+                
+            # Calculate percentage based on sizes
+            return (local_size / remote_size) * 100 if remote_size > 0 else 0
         
     except Exception as e:
         if DEBUG >= 2:
← 4c6fda7c modifying helper fucntion checking size to follow redirect f · back to Exo · comment 32560513 →