Update handler.py

handler.py

@@ -1,4 +1,5 @@
 from transformers import pipeline
+import torch
 from PIL import Image
 from io import BytesIO
 import base64
@@ -6,8 +7,18 @@ from typing import Dict, List, Any
 
 class EndpointHandler():
     def __init__(self, model_path=""):
+        # Check if a GPU is available
+        if torch.cuda.is_available():
+            device = 0
+            gpu_info = torch.cuda.get_device_name(device)
+            print(f"Using GPU: {gpu_info}")
+        else:
+            device = -1
+            cpu_info = torch.get_num_threads()
+            print(f"Using CPU with {cpu_info} threads")
+            
         # Initialize the pipeline with the specified model and set the device to GPU
-        self.pipeline = pipeline(task="zero-shot-object-detection", model=model_path, device=0)
+        self.pipeline = pipeline(task="zero-shot-object-detection", model=model_path, device=device)
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """