Use GPU if available

app.py

@@ -17,6 +17,9 @@ import easyocr
 import gradio as gr
 
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
 class MaxResize(object):
     def __init__(self, max_size=800):
         self.max_size = max_size
@@ -43,11 +46,11 @@ structure_transform = transforms.Compose([
 
 # load table detection model
 # processor = TableTransformerImageProcessor(max_size=800)
-model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-detection", revision="no_timm")
+model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-detection", revision="no_timm").to(device)
 
 # load table structure recognition model
 # structure_processor = TableTransformerImageProcessor(max_size=1000)
-structure_model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-structure-recognition-v1.1-all")
+structure_model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-structure-recognition-v1.1-all").to(device)
 
 # load EasyOCR reader
 reader = easyocr.Reader(['en'])
@@ -145,7 +148,7 @@ def visualize_detected_tables(img, det_tables):
 def detect_and_crop_table(image):
     # prepare image for the model
     # pixel_values = processor(image, return_tensors="pt").pixel_values
-    pixel_values = detection_transform(image).unsqueeze(0)
+    pixel_values = detection_transform(image).unsqueeze(0).to(device)
 
     # forward pass
     with torch.no_grad():
@@ -169,7 +172,7 @@ def detect_and_crop_table(image):
 def recognize_table(image):
     # prepare image for the model
     # pixel_values = structure_processor(images=image, return_tensors="pt").pixel_values
-    pixel_values = structure_transform(image).unsqueeze(0)
+    pixel_values = structure_transform(image).unsqueeze(0).to(device)
 
     # forward pass
     with torch.no_grad():