Update app.py

app.py

@@ -7,6 +7,7 @@ import glob
 # Ensure 'checkpoint' directory exists
 os.makedirs("checkpoint", exist_ok=True)
 
+# Function to download the model weights from a Google Drive folder
 # Function to download the model weights from a Google Drive folder
 def download_weights_from_folder(google_drive_folder_link):
     # Extract the folder ID from the Google Drive link
@@ -16,16 +17,24 @@ def download_weights_from_folder(google_drive_folder_link):
     # Download all files in the Google Drive folder
     gdown_url = f"https://drive.google.com/drive/folders/{folder_id}"
     try:
+        # Download the folder contents
         gdown.download_folder(gdown_url, quiet=False, output=output_folder)
         
-        # Check if the model file exists and rename if necessary
-        downloaded_model_path = os.path.join(output_folder, "model_state-415001.th")
-        if os.path.exists(downloaded_model_path):
-            return f"Downloaded model weights to {downloaded_model_path}"
+        # Ensure the downloaded file is named 'model_state-415001.th'
+        downloaded_files = glob.glob(os.path.join(output_folder, "*.th"))
+        if downloaded_files:
+            downloaded_model_path = downloaded_files[0]
+            target_model_path = os.path.join(output_folder, "model_state-415001.th")
+            
+            # Rename if necessary
+            if downloaded_model_path != target_model_path:
+                os.rename(downloaded_model_path, target_model_path)
+            
+            print(f"Downloaded model weights to {target_model_path}")
         else:
-            return "Model file 'model_state-415001.th' not found in the folder."
+            print("Model file 'model_state-415001.th' not found in the folder.")
     except Exception as e:
-        return f"Failed to download weights: {e}"
+        print(f"Failed to download weights: {e}")
 
 download_weights_from_folder("https://drive.google.com/drive/folders/1Bq0n-w1VT5l99CoaVg02hFpqE5eGLo9O")