spaces

app/app.py

@@ -22,6 +22,17 @@ app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
+@spaces.GPU
+def predict(filename):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return model.predict(
+        filename,   # Image to predict
+        imgsz=1024,        # Prediction image size
+        conf=0.2,          # Confidence threshold
+        device=device
+    )
+
+
 
 @app.route("/", methods=['GET', 'POST'])
 def doc_layout():
@@ -36,14 +47,7 @@ def doc_layout():
         if file and allowed_file(file.filename):
             filename = "/tmp/" + secure_filename(file.filename)
             file.save(filename)
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            print(device)
-            det_res = model.predict(
-                filename,   # Image to predict
-                imgsz=1024,        # Prediction image size
-                conf=0.2,          # Confidence threshold
-                device=device
-            )
+            det_res = predict(filename)
             names = det_res[0].names
             blocknames = [names[int(n)] for n in det_res[0].boxes.cls]
             xyxy = [a.tolist() for a in det_res[0].boxes.xyxy]

app/requirements.txt

@@ -3,5 +3,6 @@ flask == 3.1.1
 huggingface-hub == 0.33.0
 gunicorn
 --extra-index-url https://download.pytorch.org/whl/cu113
-torch == 2.5.0
+torch
+spaces