Update app.py

app.py

@@ -1,7 +1,8 @@
-# -* coding:UTF-8 -*
+# -*- coding:UTF-8 -*-
 # !/usr/bin/env python
 import numpy as np
 import gradio as gr
+import gradio.exceptions
 import roop.globals
 from roop.core import (
     start,
@@ -14,6 +15,8 @@ from PIL import Image
 import spaces
 import uuid
 import onnxruntime as ort
+import cv2
+from roop.face_analyser import get_one_face
 
 @spaces.GPU
 def swap_face(source_file, target_file, doFaceEnhancer):
@@ -32,6 +35,16 @@ def swap_face(source_file, target_file, doFaceEnhancer):
     print("source_path: ", source_path)
     print("target_path: ", target_path)
 
+    # Check if a face is detected in the source image
+    source_face = get_one_face(cv2.imread(source_path))
+    if source_face is None:
+        raise gradio.exceptions.Error("No face in source path detected.")
+
+    # Check if a face is detected in the target image
+    target_face = get_one_face(cv2.imread(target_path))
+    if target_face is None:
+        raise gradio.exceptions.Error("No face in target path detected.")
+
     output_path = os.path.join(session_dir, "output.jpg")
     normalized_output_path = normalize_output_path(source_path, target_path, output_path)
 
@@ -43,7 +56,7 @@ def swap_face(source_file, target_file, doFaceEnhancer):
     many_faces = False
     video_encoder = "libx264"
     video_quality = 18
-    max_memory = "4G"
+    max_memory = "40G"
     execution_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]  # Ưu tiên GPU, sau đó là CPU
     execution_threads = 4
     reference_face_position = 0
@@ -55,7 +68,7 @@ def swap_face(source_file, target_file, doFaceEnhancer):
     for frame_processor in get_frame_processors_modules(frame_processors):
         if not frame_processor.pre_check():
             print(f"Pre-check failed for {frame_processor}")
-            return
+            raise gradio.exceptions.Error(f"Pre-check failed for {frame_processor}")
 
     roop.globals.source_path = source_path
     roop.globals.target_path = target_path
@@ -75,7 +88,7 @@ def swap_face(source_file, target_file, doFaceEnhancer):
     roop.globals.similar_face_distance = similar_face_distance
 
     start()
-    return output_path
+    return normalized_output_path
 
 app = gr.Interface(
     fn=swap_face,