Update app.py

app.py

@@ -12,15 +12,14 @@ from Scripts.model import create_cam, create_model
 from Scripts.preprocess import crop_face, extract_face, extract_frames
 from Scripts.ca_generator import get_augs
 
-warnings.filterwarnings('ignore')
-
+import spaces
 
+warnings.filterwarnings('ignore')
 
-device = torch.device('cpu')
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 sbcl = create_model("Weights/weights.tar")
 
-
 face_detector = get_model("resnet50_2020-07-20", max_size=1024, device=device)
 face_detector.eval()
 
@@ -37,9 +36,8 @@ dlib_face_detector = dlib.get_frontal_face_detector()
 dlib_face_predictor = dlib.shape_predictor(
     'Weights/shape_predictor_81_face_landmarks.dat')
 
-
+@spaces.GPU
 def predict_image(inp):
-    
     face_list = extract_face(inp, face_detector)
 
     if len(face_list) == 0:
@@ -56,9 +54,8 @@ def predict_image(inp):
 
     return confidences, cam_image
 
-
+@spaces.GPU
 def predict_video(inp):
-    
     face_list, idx_list = extract_frames(inp, 10, face_detector)
 
     with torch.no_grad():
@@ -84,8 +81,6 @@ def predict_video(inp):
 
     return {'Real': 1-pred, 'Fake': pred}, cam_image
 
-
-
 with gr.Blocks(title="Deepfake Detection CL", theme='upsatwal/mlsc_tiet', css="""
     @import url('https://fonts.googleapis.com/css?family=Source+Code+Pro:200');
     #custom_header {
@@ -186,4 +181,4 @@ with gr.Blocks(title="Deepfake Detection CL", theme='upsatwal/mlsc_tiet', css=""
         btn_video.click(predict_video, inputs=input_video, outputs=[label_probs_video, output_image_video], api_name="/predict_video")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()