Update models.py

models.py

@@ -13,29 +13,37 @@ def download_model_from_drive(file_id, destination_path):
     """Download the model from Google Drive using gdown."""
     # Construct the Google Drive download URL
     url = f"https://drive.google.com/uc?id={file_id}"
-    # Ensure the destination directory exists
-    os.makedirs(os.path.dirname(destination_path), exist_ok=True)
+    # Extract the directory name from the destination path
+    directory = os.path.dirname(destination_path)
+    if directory:  # Ensure the directory exists only if it's valid
+        os.makedirs(directory, exist_ok=True)
     # Download the file
     gdown.download(url, destination_path, quiet=False)
 
 def load_models(device='cpu'):
     """Load YOLO model and the caption generation model."""
-    # Download the .safetensors model file if it doesn't exist
+    # Define the file path for the .safetensors model
     model_file_path = "model.safetensors"  # Adjust based on your file name
+    
+    # Download the model file if it doesn't exist
     if not os.path.exists(model_file_path):
         file_id = "1hUCqZ3X8mcM-KcwWFjcsFg7PA0hUvE3k"  # Replace with your file ID
+        print(f"Downloading model to {model_file_path}...")
         download_model_from_drive(file_id, model_file_path)
 
     # Load the YOLO model
+    print("Loading YOLO model...")
     yolo_model = YOLO("best.pt").to(device)
 
     # Load the processor for the caption model
+    print("Loading processor for the caption model...")
     processor = AutoProcessor.from_pretrained(
         "microsoft/Florence-2-base",
         trust_remote_code=True
     )
 
     # Load the caption model state dict from .safetensors
+    print("Loading caption generation model...")
     model_state_dict = load_file(model_file_path)  # Load tensors from .safetensors
     caption_model = AutoModelForCausalLM.from_pretrained(
         "microsoft/Florence-2-base",
@@ -44,6 +52,7 @@ def load_models(device='cpu'):
     caption_model.load_state_dict(model_state_dict)  # Map tensors to the model
     caption_model.to(device)  # Move the model to the correct device
 
+    print("Models loaded successfully!")
     return {
         'yolo_model': yolo_model,
         'processor': processor,
@@ -54,4 +63,4 @@ def load_models(device='cpu'):
 if __name__ == "__main__":
     device = "cuda" if torch.cuda.is_available() else "cpu"
     models = load_models(device=device)
-    print("Models loaded successfully!")
+    print("All models are ready to use!")