Update main.py

main.py

@@ -1,13 +1,14 @@
 import os
 import gdown
-from pathlib import Path
 from fastapi import FastAPI
 
 app = FastAPI()
 
 def download_weights():
-    # Create directories if they don't exist
+    # Create necessary directories if they don't exist
     os.makedirs("/tmp/.cache/gdown", exist_ok=True)
+    os.makedirs("/tmp/weights/icon_caption_florence", exist_ok=True)
+    os.makedirs("/tmp/weights/icon_detect", exist_ok=True)
 
     # Set gdown's cache directory to a writable location
     os.environ["GDOWN_CACHE"] = "/tmp/.cache/gdown"
@@ -22,8 +23,11 @@ def download_weights():
         if not os.path.exists(output_path):
             print(f"Downloading {output_path}...")
             url = f"https://drive.google.com/uc?id={file_id}"
-            gdown.download(url, output_path, quiet=False)
-            print(f"Downloaded {output_path}")
+            try:
+                gdown.download(url, output_path, quiet=False)
+                print(f"Downloaded {output_path}")
+            except Exception as e:
+                print(f"Failed to download {output_path}: {e}")
         else:
             print(f"File {output_path} already exists, skipping download")