Update app.py

app.py

@@ -15,11 +15,19 @@ MODEL_ID = os.environ.get("MODEL_ID", "BSC-LT/salamandra-7b-vision")
 DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-mtcnn = MTCNN(image_size=160, margin=0, post_process=True, device=DEVICE)
-facenet = InceptionResnetV1(pretrained='vggface2').eval().to(DEVICE)
-
 _model = None
 _processor = None
+_mtcnn = None
+_facenet = None
+
+
+def _load_face_models() -> Tuple[MTCNN, InceptionResnetV1]:
+    global _mtcnn, _facenet
+    if _mtcnn is None or _facenet is None:
+        device = DEVICE if DEVICE == "cuda" and torch.cuda.is_available() else "cpu"
+        _mtcnn = MTCNN(image_size=160, margin=0, post_process=True, device=device)
+        _facenet = InceptionResnetV1(pretrained="vggface2").eval().to(device)
+    return _mtcnn, _facenet
 
 
 def _lazy_load() -> Tuple[LlavaOnevisionForConditionalGeneration, AutoProcessor]:
@@ -96,8 +104,10 @@ def describe_batch(images: List[Image.Image], context_json: str,
     return outputs
 
 
+@spaces.GPU
 def face_image_embedding(image: Image.Image) -> List[float] | None:
     try:
+        mtcnn, facenet = _load_face_models()
         # detectar y extraer cara
         face = mtcnn(image)
 
@@ -105,7 +115,8 @@ def face_image_embedding(image: Image.Image) -> List[float] | None:
             return None
 
         # FaceNet espera tensor shape (1,3,160,160)
-        face = face.unsqueeze(0).to(DEVICE)
+        device = DEVICE if DEVICE == "cuda" and torch.cuda.is_available() else "cpu"
+        face = face.unsqueeze(0).to(device)
 
         # obtener embedding
         with torch.no_grad():