feat: improve file download reliability for Hugging Face Spaces - Add retry logic with 3 attempts for downloads - Add timeouts to prevent hanging - Add progress tracking for downloads - Improve error handling and logging - Make JSON file download more resilient

app.py

@@ -13,6 +13,7 @@ from huggingface_hub import hf_hub_download
 import easyocr
 import difflib
 import requests
+import time
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -27,18 +28,37 @@ REFERENCE_DIR = BASE_DIR / 'data' / 'planogram001'
 def download_file_from_github(url, local_path):
     """Download a file from GitHub and save it locally."""
     try:
+        # Skip if file already exists
+        if os.path.exists(local_path):
+            logger.info(f"File already exists at {local_path}, skipping download")
+            return True
+            
         # Create directory if it doesn't exist
         os.makedirs(os.path.dirname(local_path), exist_ok=True)
         
-        response = requests.get(url)
-        response.raise_for_status()
+        # Add retry logic for Hugging Face Spaces
+        max_retries = 3
+        retry_delay = 2  # seconds
         
-        with open(local_path, 'wb') as f:
-            f.write(response.content)
-        logger.info(f"Successfully downloaded {url} to {local_path}")
+        for attempt in range(max_retries):
+            try:
+                response = requests.get(url, timeout=30)  # Add timeout
+                response.raise_for_status()
+                
+                with open(local_path, 'wb') as f:
+                    f.write(response.content)
+                logger.info(f"Successfully downloaded {url} to {local_path}")
+                return True
+            except requests.exceptions.RequestException as e:
+                if attempt < max_retries - 1:
+                    logger.warning(f"Attempt {attempt + 1} failed: {str(e)}. Retrying in {retry_delay} seconds...")
+                    time.sleep(retry_delay)
+                else:
+                    raise
+                    
     except Exception as e:
         logger.error(f"Error downloading {url}: {str(e)}")
-        raise
+        return False
 
 def download_all_required_files():
     """Download all required files from GitHub repository."""
@@ -47,7 +67,17 @@ def download_all_required_files():
     # First download the JSON file to get the reference image names
     json_url = f"{base_url}/data/product_positions_adjusted_v10.json"
     json_path = "data/product_positions_adjusted_v10.json"
-    download_file_from_github(json_url, json_path)
+    
+    # Try downloading JSON file with retries
+    json_downloaded = False
+    for attempt in range(3):
+        if download_file_from_github(json_url, json_path):
+            json_downloaded = True
+            break
+        time.sleep(2)  # Wait before retry
+    
+    if not json_downloaded:
+        raise ValueError("Failed to download required JSON file after multiple attempts")
     
     # Read the JSON file to get reference image names
     global planogram_data
@@ -77,14 +107,26 @@ def download_all_required_files():
     # Create necessary directories
     os.makedirs("data/planogram001", exist_ok=True)
     
+    # Track failed downloads
+    failed_downloads = []
+    
+    # Download files with progress tracking
+    total_files = len(files_to_download)
+    downloaded_files = 0
+    
     for filename, local_path in files_to_download.items():
         url = f"{base_url}/data/{filename}"
-        try:
-            download_file_from_github(url, local_path)
-            logger.info(f"Successfully downloaded {filename} to {local_path}")
-        except Exception as e:
-            logger.error(f"Failed to download {filename}: {str(e)}")
-            raise
+        if download_file_from_github(url, local_path):
+            downloaded_files += 1
+            logger.info(f"Progress: {downloaded_files}/{total_files} files downloaded")
+        else:
+            failed_downloads.append(filename)
+    
+    if failed_downloads:
+        logger.error(f"Failed to download the following files: {', '.join(failed_downloads)}")
+        raise ValueError(f"Failed to download {len(failed_downloads)} required files")
+    
+    logger.info("Successfully downloaded all required files")
 
 # Initialize CLIP model and processor with a smaller model
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -167,9 +209,6 @@ except Exception as e:
 def load_reference_images():
     """Load all reference product images and pre-compute their features."""
     try:
-        # First ensure all required files are downloaded
-        download_all_required_files()
-        
         # Now load the reference images
         for img_path in REFERENCE_DIR.glob('*.png'):
             if img_path.name == 'empty-space.png' or img_path.name == 'planogram.png':
@@ -292,18 +331,22 @@ def compare_images(ref_img, uploaded_img):
 @cl.on_chat_start
 async def start():
     """Initialize the chat session."""
-    # Load reference images
-    load_reference_images()
-    
-    ref_img = load_reference_image()
-    if ref_img is None:
-        await cl.Message(
-            content="Error: Reference planogram image not found. Please ensure planogram.png exists in the project directory."
-        ).send()
-        return
-    
-    # Create a welcome message with instructions
-    welcome_msg = """
+    try:
+        # First ensure all required files are downloaded
+        download_all_required_files()
+        
+        # Then load the reference images
+        load_reference_images()
+        
+        ref_img = load_reference_image()
+        if ref_img is None:
+            await cl.Message(
+                content="Error: Reference planogram image not found. Please ensure planogram.png exists in the project directory."
+            ).send()
+            return
+        
+        # Create a welcome message with instructions
+        welcome_msg = """
 # Welcome to Planolyzer! 🛍️
 
 ## Quick Start:
@@ -313,61 +356,67 @@ async def start():
    - Adding empty spaces to the reference planogram
    - Moving items to wrong positions
 """
-    
-    # Send welcome message
-    await cl.Message(content=welcome_msg).send()
-    
-    # Send reference planogram
-    await cl.Message(
-        content="## Reference Planogram:",
-        elements=[
-            cl.Image(
-                name="planogram",
-                path=str(REFERENCE_DIR / 'planogram.png'),
-                display="inline",
-                size="medium"
-            ),
-            cl.File(
-                name="planogram.png",
-                path=str(REFERENCE_DIR / 'planogram.png'),
-                display="inline"
-            ),
-            cl.Image(
-                name="empty_space",
-                path=str(REFERENCE_DIR / 'empty-space.png'),
-                display="inline",
-                size="small"
-            ),
-            cl.File(
-                name="empty-space.png",
-                path=str(REFERENCE_DIR / 'empty-space.png'),
-                display="inline"
-            )
-        ]
-    ).send()
-    
-    # Send test image
-    await cl.Message(
-        content="## Test Image:",
-        elements=[
-            cl.Image(
-                name="test_shelf",
-                path=str(BASE_DIR / 'data' / 'test_shelf_image_cig_003.png'),
-                display="inline",
-                size="medium"
-            ),
-            cl.File(
-                name="test_shelf_image_cig_003.png",
-                path=str(BASE_DIR / 'data' / 'test_shelf_image_cig_003.png'),
-                display="inline"
-            )
-        ]
-    ).send()
-    
-    # Send additional instructions
-    await cl.Message(
-        content="Try downloading and uploading the test image first to see how the system works!"
-    ).send()
+        
+        # Send welcome message
+        await cl.Message(content=welcome_msg).send()
+        
+        # Send reference planogram
+        await cl.Message(
+            content="## Reference Planogram:",
+            elements=[
+                cl.Image(
+                    name="planogram",
+                    path=str(REFERENCE_DIR / 'planogram.png'),
+                    display="inline",
+                    size="medium"
+                ),
+                cl.File(
+                    name="planogram.png",
+                    path=str(REFERENCE_DIR / 'planogram.png'),
+                    display="inline"
+                ),
+                cl.Image(
+                    name="empty_space",
+                    path=str(REFERENCE_DIR / 'empty-space.png'),
+                    display="inline",
+                    size="small"
+                ),
+                cl.File(
+                    name="empty-space.png",
+                    path=str(REFERENCE_DIR / 'empty-space.png'),
+                    display="inline"
+                )
+            ]
+        ).send()
+        
+        # Send test image
+        await cl.Message(
+            content="## Test Image:",
+            elements=[
+                cl.Image(
+                    name="test_shelf",
+                    path=str(BASE_DIR / 'data' / 'test_shelf_image_cig_003.png'),
+                    display="inline",
+                    size="medium"
+                ),
+                cl.File(
+                    name="test_shelf_image_cig_003.png",
+                    path=str(BASE_DIR / 'data' / 'test_shelf_image_cig_003.png'),
+                    display="inline"
+                )
+            ]
+        ).send()
+        
+        # Send additional instructions
+        await cl.Message(
+            content="Try downloading and uploading the test image first to see how the system works!"
+        ).send()
+
+    except Exception as e:
+        logger.error(f"Error initializing chat session: {str(e)}")
+        await cl.Message(
+            content=f"Error initializing chat session: {str(e)}"
+        ).send()
 
 @cl.on_message
 async def main(message: cl.Message):