Update models.py

models.py

@@ -15,24 +15,29 @@ def load_models(device='cpu'):
     # Set default dtype for torch
     torch.set_default_dtype(torch.float32)
 
-    # Download the model from Google Drive (if not already present)
-    model_file_path = 'model.safetensors'  # Use the correct model file name
+    # Download the caption model if not present
+    model_file_path = 'model.safetensors'  # Adjust this name as needed
     if not os.path.exists(model_file_path):
         file_id = "1hUCqZ3X8mcM-KcwWFjcsFg7PA0hUvE3k"  # Replace with your Google Drive file ID
         download_model_from_drive(file_id, model_file_path)
-    
+
     # Load the YOLO model
     yolo_model = YOLO('best.pt').to(device)
 
-    # Load processor and caption model
+    # Load processor for the caption model
     processor = AutoProcessor.from_pretrained(
         "microsoft/Florence-2-base", 
         trust_remote_code=True
     )
 
-    # Load the caption model from the downloaded .safetensors file
-    # Use safetensors library to load the model
-    caption_model = load_file(model_file_path, framework="pt", device=device)
+    # Load the caption model
+    model_state_dict = load_file(model_file_path)  # Load tensors from .safetensors
+    caption_model = AutoModelForCausalLM.from_pretrained(
+        "microsoft/Florence-2-base", 
+        trust_remote_code=True
+    )
+    caption_model.load_state_dict(model_state_dict)  # Map tensors to the model
+    caption_model.to(device)  # Move the model to the correct device
     
     return {
         'yolo_model': yolo_model,