Update

app.py

@@ -2,13 +2,13 @@
 
 from __future__ import annotations
 
-import os
 import pathlib
 import sys
 
 import cv2
 import gradio as gr
 import numpy as np
+import spaces
 import torch
 
 sys.path.insert(0, "face_detection")
@@ -20,15 +20,29 @@ from ibug.face_detection import RetinaFacePredictor
 DESCRIPTION = "# [ibug-group/face_alignment](https://github.com/ibug-group/face_alignment)"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-detector = RetinaFacePredictor(threshold=0.8, device=device, model=RetinaFacePredictor.get_model("mobilenet0.25"))
+detector = RetinaFacePredictor(threshold=0.8, device="cpu", model=RetinaFacePredictor.get_model("mobilenet0.25"))
+detector.device = device
+detector.net.to(device)
+
+
+def load_model(model_name: str, device: torch.device) -> FANPredictor:
+    model = FANPredictor(
+        device="cpu", model=FANPredictor.get_model(model_name), config=FANPredictor.create_config(use_jit=False)
+    )
+    model.device = device
+    model.net.to(device)
+    return model
+
+
 model_names = [
     "2dfan2",
     "2dfan4",
     "2dfan2_alt",
 ]
-models = {name: FANPredictor(device=device, model=FANPredictor.get_model(name)) for name in model_names}
+models = {name: load_model(name, device) for name in model_names}
 
 
+@spaces.GPU
 def predict(image: np.ndarray, model_name: str, max_num_faces: int, landmark_score_threshold: int) -> np.ndarray:
     model = models[model_name]
 
@@ -72,7 +86,6 @@ with gr.Blocks(css="style.css") as demo:
         inputs=[image, model_name, max_num_faces, landmark_score_thrshold],
         outputs=result,
         fn=predict,
-        cache_examples=os.getenv("CACHE_EXAMPLES") == "1",
     )
     run_button.click(
         fn=predict,