go back to GPU

app.py

@@ -10,23 +10,22 @@ from utils.florence import load_florence_model, run_florence_inference, \
     FLORENCE_OPEN_VOCABULARY_DETECTION_TASK
 from utils.sam import load_sam_image_model, run_sam_inference
 
-# DEVICE = torch.device("cuda")
-DEVICE = torch.device("cpu")
+DEVICE = torch.device("cuda")
+# DEVICE = torch.device("cpu")
 
-# torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
-# if torch.cuda.get_device_properties(0).major >= 8:
-#     torch.backends.cuda.matmul.allow_tf32 = True
-#     torch.backends.cudnn.allow_tf32 = True
+torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
+if torch.cuda.get_device_properties(0).major >= 8:
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
 
 
 FLORENCE_MODEL, FLORENCE_PROCESSOR = load_florence_model(device=DEVICE)
 SAM_IMAGE_MODEL = load_sam_image_model(device=DEVICE)
 
 
-# @spaces.GPU
-# @torch.inference_mode()
-# @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+@spaces.GPU
 @torch.inference_mode()
+@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def process_image(image_input, text_input) -> Optional[Image.Image]:
     if not image_input:
         gr.Info("Please upload an image.")