Update app.py

app.py

@@ -41,28 +41,46 @@ class Predictor:
             print("No face found")
             return None
 
+    # 这是修正后的predict函数，请用它替换掉您源代码中的旧版本
     def predict(self, input_image, swap_image):
         """Run a single prediction on the model"""
+        # 【新增】检查输入是否为空
+        if input_image is None or swap_image is None:
+            raise gr.Error("请确保同时上传了目标图片和源图片。") # 使用Gradio的方式优雅地报错
+    
         try:
-            frame = cv2.imread(input_image.name)
-            face = self.get_face(frame)
-            source_face = self.get_face(cv2.imread(swap_image.name))
-            try:
-                print(frame.shape, face.shape, source_face.shape)
-            except:
-                print("printing shapes failed.")
-            result = self.face_swapper.get(frame, face, source_face, paste_back=True)
-
+            # 读取图片
+            target_img = cv2.imread(input_image.name)
+            swap_img = cv2.imread(swap_image.name)
+    
+            # 分析人脸
+            target_face = self.get_face(target_img)
+            source_face = self.get_face(swap_img)
+    
+            # 【关键修正】在执行换脸前，检查是否成功找到了人脸
+            if target_face is None:
+                raise gr.Error("在目标图片中未能检测到人脸，请更换图片后重试。")
+            if source_face is None:
+                raise gr.Error("在源图片中未能检测到人脸，请更换图片后重试。")
+            
+            # 如果人脸都找到了，才执行核心换脸操作
+            result = self.face_swapper.get(target_img, target_face, source_face, paste_back=True)
+    
+            # 增强画质
             _, _, result = self.face_enhancer.enhance(
                 result,
                 paste_back=True
             )
+    
+            # 保存并返回结果
             out_path = tempfile.mkdtemp() + f"/{str(int(time.time()))}.jpg"
             cv2.imwrite(out_path, result)
             return out_path
+        
         except Exception as e:
-            print(f"{e}")
-            return None
+            # 如果发生其他未知错误，也通过Gradio报错
+            print(f"An unexpected error occurred: {e}")
+            raise gr.Error(f"发生未知错误: {e}")
 
 
 # Instantiate the Predictor class