Update app.py

app.py

@@ -12,8 +12,10 @@ import numpy as np
 from PIL import Image
 import spaces
 
+
+device = torch.device( cuda if torch.cuda.is_available() else cpu )
 processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble")
-model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble").to("cuda")
+model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble").to(device)
 
 BOUNDING_BOX_ANNOTATOR = sv.BoundingBoxAnnotator()
 MASK_ANNOTATOR = sv.MaskAnnotator()
@@ -76,7 +78,7 @@ def process_video(
 
 @spaces.GPU
 def query(image, texts):
-  inputs = processor(text=texts, images=image, return_tensors="pt").to("cuda")
+  inputs = processor(text=texts, images=image, return_tensors="pt").to(device)
   with torch.no_grad():
     outputs = model(**inputs)
   target_sizes = torch.Tensor([image.shape[:-1]])